diff --git a/.agents/.gitignore b/.agents/.gitignore new file mode 100644 index 00000000..3cecb244 --- /dev/null +++ b/.agents/.gitignore @@ -0,0 +1,7 @@ +*.md +!CLAUDE.md + +# Shared, versioned repository skills. +!skills/ +!skills/**/ +!skills/**/SKILL.md diff --git a/.agents/skills/postgkyl-api-cli/SKILL.md b/.agents/skills/postgkyl-api-cli/SKILL.md new file mode 100644 index 00000000..77607366 --- /dev/null +++ b/.agents/skills/postgkyl-api-cli/SKILL.md @@ -0,0 +1,39 @@ +--- +name: postgkyl-api-cli +description: Add or change public Postgkyl verbs, signatures, fluent aliases, or generated CLI behavior while preserving API/CLI parity. +--- + +# Keep one executable public API + +Core operations use `op(data: GDataState, *, ..., inplace=False, tag=None, +label=None) -> GDataState`; funnel mutation/new-result decisions through `_result`. + +- Alias exact fluent verbs in the GData class body to canonical operations; no + wrappers, runtime setattr, or lazy imports. GDataGroup broadcasts member verbs. + Multi-dataset verbs share the aliases in `gdata/verbs.py`. +- Rendering lives in `render/`. Facade, operations, fluent methods, and generated + CLI refer to the same canonical callables. +- Every public boolean defaults to `False`. For enabled default behavior, name + its inverse (`no_show`, `volume`, `nodal`) and implement it under the false + branch. CLI bare boolean options mean True; explicit True/False remain accepted. +- Expose basis/order/value_form overrides only at load time, never downstream. + +The CLI discovers the public API, constructs CommandSpec records, and lowers them +through one generic compiler. Keep scientific command knowledge in API signatures, +annotations, docstrings, and command metadata. Do not add handwritten subcommands +or a `cli/commands/` package. + +Preserve underscores exactly (`local_poly`, `--num_moms`). Assign a short option +to the first parameter with each initial in signature order; reserve `-h` for +help. Aliases/abbreviations change spelling only. Bare filenames expand to +`load --file_name`; they do not define different loading semantics. + +Use native Click chaining and callback-before-dispatch. Help groups the flat +inventory into Verbs, Diagnostics, Render, and Utility. The console entry point +is `postgkyl.cli.app:cli`. Version reporting is owned by `_version.py`, exported +through the facade, and receives the version string explicitly to avoid an import +ordering dependency; preserve its build/dependency diagnostics. + +Verify signature changes with `tests/test_cli_generator.py`, +`tests/test_cli_commands.py`, and relevant diagnostic CLI tests; check canonical +callable identities in `tests/test_postgkyl.py` when changing aliases. diff --git a/.agents/skills/postgkyl-architecture/SKILL.md b/.agents/skills/postgkyl-architecture/SKILL.md new file mode 100644 index 00000000..17462d09 --- /dev/null +++ b/.agents/skills/postgkyl-architecture/SKILL.md @@ -0,0 +1,39 @@ +--- +name: postgkyl-architecture +description: Place modules and enforce import boundaries when adding or refactoring Postgkyl operations, diagnostics, rendering, or public surfaces. +--- + +# Preserve layer ownership + +`tests/test_postgkyl.py::_ALLOWED` and its AST checks are the authority for +allowed imports and acyclicity. Consult that mapping for exact edges; when +adding a top-level module or allowed edge, update the contract and this guidance +in the same change if ownership changes. Do not maintain a second import matrix. + +| Layer | Owns | +| --- | --- | +| `__init__.py`, `cli/`, `gui/` | Public surfaces; the facade only re-exports, with no function/class definitions | +| `diagnostics/` | Equation-specific physical conclusions and compositions | +| `gdata/` | Fluent GData/GDataGroup API and load entry point | +| `operations/` | Data transformations; flat equation-blind verbs | +| `render/` | Canonical plotting and animation callables | +| `gdatastate/` | State, clone, `_result`, shared guards and point-value materialization; no verbs | +| `dg/`, `io/` | Kernel orchestration and file I/O respectively | +| `numerics/`, `cli_spec.py` | Pure math and frozen command metadata; no internal imports | +| `gpython/` | Sole foreign-library boundary | + +Imports point downward along allowed edges, never back to a higher surface. +Operations accept GDataState; `_result` constructs `type(self)` so fluent results +remain GData without importing gdata. Reuse `gdatastate/guards.py` and +`materialize.py` instead of duplicating capability checks or native bridges. +Readers return `(grid, values)` and fill plain metadata; they never import state. + +Diagnostics are free functions under model families `gk`, `vm`, `pkpm`, or `mom`, +not GData methods. Model-specific loading belongs beside its physics. Resolve +output stems/frames through `diagnostics/discovery.py`; keep quantity vocabulary +in the equation module's `VARIABLES` table. Use public functions from lower +layers and return a state via `_result`, or a Figure for program diagnostics. +for the current major version. + +Run the import, foreign-floor, facade, and canonical-callable contracts in +`tests/test_postgkyl.py` after structural changes. diff --git a/.agents/skills/postgkyl-data/SKILL.md b/.agents/skills/postgkyl-data/SKILL.md new file mode 100644 index 00000000..a97bc9d0 --- /dev/null +++ b/.agents/skills/postgkyl-data/SKILL.md @@ -0,0 +1,43 @@ +--- +name: postgkyl-data +description: Preserve DG representation and backend semantics when changing loading, arithmetic, conversions, integration, or terminal data consumers. +--- + +# Preserve data meaning + +Storage (`backend`: `gkyl` or `numpy`) and representation +(`ctx["value_form"]`: `modal`, `nodal`, or `quad`) are separate facts. Do not +add an `is_modal` flag or infer point-value capabilities from storage alone. + +- Modal coefficients use Gkeyll DG operations: weak multiply/divide, coefficient + linear combinations, scalar scaling/mean shifts, repeated weak multiplication + for integer powers, and native integration. Ufuncs, array conversion, selection, + and plotting must not treat coefficients as field values. +- Nodal and quadrature values are fields at points. Pointwise NumPy operations + are allowed; native results wrap back into native storage in the same value_form. + Plot true point locations; non-tensor node sets use explicit `.to_quad()`. +- Reject mixed backends or mixed value_forms when combining datasets. +- `.interpolate()` is the one-way bridge to a new, by-value NumPy field array. + Interpolation matrices come from Gkeyll basis functions, applied per cell; + nodal input first uses the exact nodal-to-modal transform. +- Only explicit `.to_modal()`, `.to_nodal()`, and `.to_quad()` perform representation + conversions. `.apply(fn, num_quad=...)` spells modal → quad → fn → projection, + equivalent to `fn(d.to_quad()).to_modal()`. Respect quadrature exactness limits. +- Full modal integration is terminal and native; partial integration uses native + averaging with physical-volume scaling and returns lower-dimensional modal data. + `.average()` and `.eval_at_coord_proj()` likewise stay modal/native and compose. + `.local_poly()` builds a discontinuity-preserving plotting mesh. + +Resolve `basis_type`, `poly_order`, and `value_form` once from file metadata or +explicit load options. Downstream verbs read ctx and raise if required metadata +is missing; never add basis/order override parameters to them. + +For spatial data with unresolved basis metadata, loading warns and defaults to +serendipity p0 nodal data. When value_form is defaulted, express the grid as cell +centers, matching one point per cell. Dynvectors have no spatial DG basis and +are exempt. The readers' narrower case—known basis but no value_form tag—assumes +modal silently; keep these two defaults distinct. + +Route results through `_result` and terminal native point-value consumers through +`gdatastate.materialize_point_values`. Preserve native ownership and read-only +coefficient views; use the existing representation tests to verify semantics. diff --git a/.agents/skills/postgkyl-design/SKILL.md b/.agents/skills/postgkyl-design/SKILL.md new file mode 100644 index 00000000..362e39b1 --- /dev/null +++ b/.agents/skills/postgkyl-design/SKILL.md @@ -0,0 +1,82 @@ +--- +name: postgkyl-design +description: Apply Postgkyl coding doctrine when designing, implementing, or reviewing Python changes. +--- + +# Coding Doctrine + +**0. Locality of reasoning.** Every principle below is a projection of +one axiom: a reader must be able to understand a fragment without the +whole program. Whatever keeps a local conclusion sound — a frozen +record, an honest signature, a stated law — is doctrine. Whatever +forces a global search — ambient state, a leaky layer, a second copy +of a fact — is the enemy. + +*Data — what it does, and what it may say* + +**I. Data is inert. Functions transform.** No objects that know +things and do things. Data is a frozen record. Behavior is a function +that takes data in and returns data out. If you're reaching for +inheritance, you've taken a wrong turn. + +**II. Make illegal states unrepresentable.** The shape of a datum is +its strongest invariant. Constructors refuse invalid states; a checked +fact becomes a type; downstream never re-proves what upstream +established. Parse, don't validate. + +*Functions — one idea, honestly declared* + +**III. A function is one idea.** It takes exactly what it needs and +returns exactly what it computes. If the signature has two concepts in +it, you have two functions. + +**IV. The signature tells the whole truth.** Inward: if something +needs a value, it receives it as a parameter — no spooky action at a +distance, no stringly-typed interfaces, no implicit state. Outward: +same inputs, same outputs; effects and failure appear in the type, not +in the fine print. Pure core, effects at the edges. + +*Knowledge — one home per fact* + +**V. Every fact has one home.** One authoritative representation of +each decision and each piece of knowledge; everything else inherits or +is derived mechanically — never maintained by hand in parallel. +Configuration is decided once, at the highest level, and threaded +down; no module ever decides its own context. If the design and the +implementation can disagree, you have two sources of truth and zero. + +*Layers — what above, how below* + +**VI. Separate what from how.** Logic and machinery are different +concerns with a hard boundary. The layer that says *what* to compute +should be readable by someone who has never seen the machinery +underneath. The layer that says *how* lives below, stays below, and +nothing leaks up from it. + +**VII. Notation is execution; lowering is transliteration.** Looking +up: the spec layer reads like the math or logic it implements — when +notation *is* the executable object, not a comment beside it, bugs +have nowhere to hide. Looking down: the layer that executes the spec +reproduces it exactly — nothing added, nothing dropped, nothing +reinterpreted; no opinions, no defaults, no helpful conversions. If +the lowering changes anything, the spec is a lie. + +*Abstraction — earned, and binding* + +**VIII. Earn your abstractions.** No abstraction before the second +use. Three similar lines is better than a premature helper. The right +amount of complexity is the minimum the current task demands — not the +current task plus three hypothetical future ones. + +**IX. An abstraction is a contract.** It is defined by what it +guarantees, not what it hides. If you can't state what is always true +of it — properties a client may rely on without reading the +implementation — it isn't an abstraction, it's indirection. Two +implementations that honor the contract must be interchangeable; and +its outputs stay in its vocabulary, so uses compose. + +*Verification — formal first* + +**X. Trust the most formal thing first.** Types over tests, tests +over docs, docs over comments. Invest in whichever layer catches the +bug earliest with the least ongoing maintenance cost. diff --git a/.agents/skills/postgkyl-development/SKILL.md b/.agents/skills/postgkyl-development/SKILL.md new file mode 100644 index 00000000..dd1d505c --- /dev/null +++ b/.agents/skills/postgkyl-development/SKILL.md @@ -0,0 +1,48 @@ +--- +name: postgkyl-development +description: Set up Postgkyl development, choose verification commands, or maintain shared agent configuration and Entire hooks. +--- + +# Develop and verify + +Install NumPy before the extension so it builds against the runtime ABI: + +```bash +pip install --upgrade numpy setuptools wheel +pip install --no-build-isolation -e '.[test]' +``` + +Run focused tests for the change, then the required broader checks: + +```bash +pytest tests/ +# Without an editable install: +PYTHONPATH=src python -m pytest tests/ +``` + +Useful API/CLI smoke checks (replace data filenames with suitable fixtures): + +```bash +pgkyl --help +pgkyl --version +pgkyl file.gkyl info +pgkyl file.gkyl interpolate select --z0 0 plot +pgkyl euler_5m_0.gkyl interpolate five_moment_pressure --num_moms 5 plot +pgkyl a.gkyl b.gkyl evaluate "f0 f1 +" interpolate plot +``` + +Shared agent configuration lives in `.agents/`; `.claude` and `.codex` are +relative symlinks to it. Root `CLAUDE.md` links to `AGENTS.md`. Put focused skills +in `.agents/skills//SKILL.md` with name and description frontmatter. Keep +root instructions short and route only to skills relevant to the task. + +Preserve Entire's Claude hooks in `settings.json` and Codex hooks in `hooks.json` +inside the shared directory, including matchers, commands, and timeouts. Their +original `.claude/settings.json` and `.codex/hooks.json` paths must still resolve. +Leave `.entire/` and Entire's Git hooks intact. Retain existing agent definitions +and historical plans when reorganizing folders; historical markdown remains +ignored, while new skill entry points must be visible to Git. + +For instruction/configuration-only changes, validate skill frontmatter, relative +links, Git visibility, and hook preservation; running scientific tests is not +necessary unless application behavior changes. diff --git a/.agents/skills/postgkyl-native/SKILL.md b/.agents/skills/postgkyl-native/SKILL.md new file mode 100644 index 00000000..d3a89df2 --- /dev/null +++ b/.agents/skills/postgkyl-native/SKILL.md @@ -0,0 +1,35 @@ +--- +name: postgkyl-native +description: Change or troubleshoot the Gkeyll native bridge, basis matrices, readers, memory ownership, or native builds in Postgkyl. +--- + +# Keep the foreign boundary compiled + +`gpython/` is the only doorway to Gkeyll. No ctypes declarations, native struct +layouts, or foreign signatures in Python or other layers. `gpython.available()` +is the single capability switch. + +The shim lives in `gkeyll/core/zero/{gkyl_gpython.h,gpython.c}` and builds into +Gkeyll's `libg0core.so`. Struct access, by-value basis conventions, and function +pointer dispatch belong there, checked against that tree's headers. +`gpython/csrc/_gpythonmodule.c` wraps only the shim's opaque handles, scalars, +and buffers. `_lib.py` imports the extension and checks `GPYTHON_API_VERSION`. + +Build against the pinned clone using `scripts/build_gkeyll.sh` and +`scripts/build_gpython.sh`. Bundle libg0core beside the extension and retain +relative `$ORIGIN`/`@loader_path` linking. Preserve generated build provenance; +Gkeyll is a build-time clone and need not exist at runtime. + +GkylArray capsules own and release arrays. Zero-copy construction pins its NumPy +buffer; view base chains pin the capsule so views outlive their dataset. Never +return unowned C memory. Build interpolation and representation matrices by +evaluating Gkeyll's own basis through the shim, not duplicated basis formulas. + +`dg/` orchestrates kernels; `io/` dispatches readers. Prefer GkylCReader for native +field reads; retain the Python reader fallback for unavailable native libraries, +partial loads, and dynvectors. Keep NumPy installed before building so the +extension uses the runtime ABI; use `pip install --no-build-isolation -e '.[test]'`. + +Verify handshake, memory lifetime, basis/interpolation, and modal algebra using +the relevant existing tests. Report native test skips explicitly when no compiled +library is available; skips do not establish native correctness. diff --git a/.agents/skills/postgkyl-testing/SKILL.md b/.agents/skills/postgkyl-testing/SKILL.md new file mode 100644 index 00000000..7643fd5d --- /dev/null +++ b/.agents/skills/postgkyl-testing/SKILL.md @@ -0,0 +1,12 @@ +--- +name: postgkyl-testing +description: Design and best practices for testing Postgkyl code, including unit tests and examples. +--- + +# Testing + +Unit tests aim for 100% code coverage where possible. Every change must have an associated unit test. Every bug fix must have a test. Every feature must be tested. Run `pytest` to check that all tests pass. + +Data for testing must be generated automatically by `tests/generate_test_data.py`. This keeps the repository light. DO NOT COMMIT LARGE DATA FILES TO GIT. + +Examples are real-world use cases of postgkyl. They read data and make plots. Examples are displayed in the documentation. diff --git a/.clang-format b/.clang-format new file mode 100644 index 00000000..0a3eab5b --- /dev/null +++ b/.clang-format @@ -0,0 +1,21 @@ +--- +Language: Cpp +BasedOnStyle: LLVM + +ColumnLimit: 80 + +IndentWidth: 2 +ContinuationIndentWidth: 4 +TabWidth: 2 +UseTab: Never + +AlwaysBreakAfterReturnType: AllDefinitions +BreakBeforeBraces: WebKit +AllowShortFunctionsOnASingleLine: None + +DerivePointerAlignment: false +PointerAlignment: Right + +IncludeBlocks: Preserve +SortIncludes: CaseSensitive +... diff --git a/.claude b/.claude new file mode 120000 index 00000000..c0ca4685 --- /dev/null +++ b/.claude @@ -0,0 +1 @@ +.agents \ No newline at end of file diff --git a/.codex b/.codex new file mode 120000 index 00000000..c0ca4685 --- /dev/null +++ b/.codex @@ -0,0 +1 @@ +.agents \ No newline at end of file diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..8c1b62c6 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,22 @@ +version: 2 + +updates: + - package-ecosystem: pip + directory: / + schedule: + interval: weekly + groups: + python-dependencies: + patterns: + - "*" + open-pull-requests-limit: 5 + + - package-ecosystem: github-actions + directory: / + schedule: + interval: weekly + groups: + actions: + patterns: + - "*" + open-pull-requests-limit: 5 diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 00000000..fe68798a --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,50 @@ +name: Documentation + +on: + pull_request: + push: + branches: [main] + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: docs-${{ github.ref }} + cancel-in-progress: true + +jobs: + documentation: + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + MPLBACKEND: Agg + POSTGKYL_REQUIRE_GKEYLL: "1" + VTK_DEFAULT_OPENGL_WINDOW: vtkEGLRenderWindow + LIBGL_ALWAYS_SOFTWARE: "1" + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + - name: Install native Postgkyl and documentation tools + run: | + sudo apt-get update + sudo apt-get install -y libegl1 libgl1-mesa-dri + python -m pip install --upgrade numpy setuptools wheel + python -m pip install --no-build-isolation -e '.[docs,test]' + - name: Test documentation contracts and downloadable examples + run: python -m pytest tests/test_documentation.py tests/test_examples.py tests/test_docs_build.py + - name: Build a reviewable website + run: | + python scripts/build_docs.py + python -m sphinx -W --keep-going -b html -c docs build/docs/source build/docs/html + - uses: actions/upload-artifact@v4 + with: + name: postgkyl-docs + path: build/docs/html + if-no-files-found: error diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 8c18e600..e85863ff 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,36 +1,337 @@ - -name: Test Postgkyl +name: CI on: pull_request: - types: - - opened - - reopened - - synchronize - - ready_for_review + types: [opened, reopened, synchronize, ready_for_review] + push: + branches: [main] + workflow_dispatch: + schedule: + # Exercise external Chrome/ffmpeg integrations even during quiet weeks. + - cron: "0 7 * * 1" + +# Every job is read-only. In particular, formatting is a check on the source +# commit, never a mutation which is passed on to later jobs. +permissions: + contents: read + +concurrency: + group: ci-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +env: + # pull_request checks run the contributor's exact commit, rather than a + # synthetic merge commit. The fallbacks cover push and workflow_dispatch. + SOURCE_REPOSITORY: ${{ github.event.pull_request.head.repo.full_name || github.repository }} + SOURCE_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + FORCE_COLOR: "1" + PYTHONFAULTHANDLER: "1" + PYTHONUNBUFFERED: "1" jobs: - install_and_test: + quality: + name: Quality (format and lint) + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Check out source commit + uses: actions/checkout@v4 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + - name: Verify source commit + shell: bash + run: test "$(git rev-parse HEAD)" = "$SOURCE_SHA" + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + - name: Install quality tools + run: python -m pip install pre-commit==4.3.0 ruff==0.12.11 + - name: Validate hooks and check formatting + run: | + pre-commit validate-config + pre-commit run --all-files --show-diff-on-failure + - name: Lint + run: ruff check . + compatibility: + name: Core / Python ${{ matrix.python-version }} runs-on: ubuntu-latest + timeout-minutes: 45 + env: + POSTGKYL_SKIP_GKEYLL_BUILD: "1" strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12"] + # The authoritative full suite runs on 3.12 below. This matrix covers + # the other supported interpreters without repeating costly renderer + # and native integration tests. + include: + - python-version: "3.10" + numpy: "numpy==2.2.6" + - python-version: "3.11" + numpy: "numpy==2.2.6" + - python-version: "3.13" + numpy: "numpy==2.2.6" + # NumPy 2.2 predates Python 3.14 support. Resolve the current NumPy + # release on this one lane instead of asking pip for no wheel. + - python-version: "3.14" + numpy: "numpy" + steps: + - name: Check out source commit + uses: actions/checkout@v4 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + - name: Verify source commit + shell: bash + run: test "$(git rev-parse HEAD)" = "$SOURCE_SHA" + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: pyproject.toml + - name: Install Postgkyl and test dependencies + run: | + python -m pip install "${{ matrix.numpy }}" setuptools wheel + python -m pip install --no-build-isolation -e ".[test]" + - name: Run core tests + run: >- + python -m pytest + -m compatibility + --strict-config --strict-markers + --timeout=120 + coverage: + name: Full suite and branch coverage / Linux + runs-on: ubuntu-latest + timeout-minutes: 60 + env: + MPLBACKEND: Agg + POSTGKYL_REQUIRE_GKEYLL: "1" + steps: + - name: Check out source commit + uses: actions/checkout@v4 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + - name: Verify source commit + shell: bash + run: test "$(git rev-parse HEAD)" = "$SOURCE_SHA" + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + - name: Install Postgkyl and test dependencies + run: | + python -m pip install "numpy==2.2.6" setuptools wheel + python -m pip install --no-build-isolation -e ".[test]" + - name: Start headless display + uses: pyvista/setup-headless-display-action@v4 + - name: Require the compiled Gkeyll bridge + run: | + python - <<'PY' + from postgkyl import gpython + gpython.require() + print(gpython.lib_path()) + PY + - name: Run full suite with coverage + run: >- + python -m pytest + -m "not external_tool" + --strict-config --strict-markers + --timeout=120 + --cov=postgkyl --cov-branch + --cov-report=term-missing + --cov-report=xml:coverage.xml + --cov-fail-under=99 + - name: Upload coverage report + if: always() + uses: actions/upload-artifact@v4 + with: + name: coverage-xml + path: coverage.xml + if-no-files-found: warn + retention-days: 14 + + native_macos: + name: Native bridge / macOS + runs-on: macos-latest + timeout-minutes: 60 + env: + POSTGKYL_REQUIRE_GKEYLL: "1" + steps: + - name: Check out source commit + uses: actions/checkout@v4 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + - name: Verify source commit + shell: bash + run: test "$(git rev-parse HEAD)" = "$SOURCE_SHA" + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + - name: Install Postgkyl and test dependencies + run: | + python -m pip install "numpy==2.2.6" setuptools wheel + python -m pip install --no-build-isolation -e ".[test]" + - name: Require the compiled Gkeyll bridge + run: | + python - <<'PY' + from postgkyl import gpython + gpython.require() + print(gpython.lib_path()) + PY + - name: Run native tests in isolated processes + # Native-library aborts on macOS should fail one test, not terminate + # the entire session and hide which test caused the failure. + # pytest-forked already captures each child's stdout/stderr. A second + # capture layer makes pytest replace those streams during setup; on + # macOS they can then be finalized unclosed, and warnings-as-errors + # turns the resulting ResourceWarning into an error for every test. + run: >- + python -m pytest + -m native + --strict-config --strict-markers + --timeout=120 + --capture=no + --forked + --ignore=tests/test_examples.py + - name: Run tutorial tests in a fresh interpreter + # PyVista rendering initializes macOS Objective-C classes, which is + # unsafe in pytest-forked's fork-without-exec children. Keep these + # examples in a separate, unforked pytest session so they retain the + # GUI guards and VTK shutdown handling in tests/conftest.py. + run: >- + python -m pytest tests/test_examples.py + --strict-config --strict-markers + --timeout=120 + + renderers: + name: Renderer integrations / Linux + runs-on: ubuntu-latest + timeout-minutes: 45 + env: + MPLBACKEND: Agg + POSTGKYL_SKIP_GKEYLL_BUILD: "1" + steps: + - name: Check out source commit + uses: actions/checkout@v4 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + - name: Verify source commit + shell: bash + run: test "$(git rev-parse HEAD)" = "$SOURCE_SHA" + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + - name: Set up Chrome for Kaleido + if: github.event_name != 'pull_request' + uses: browser-actions/setup-chrome@v2 + - name: Install Postgkyl and test dependencies + run: | + python -m pip install "numpy==2.2.6" setuptools wheel + python -m pip install --no-build-isolation -e ".[test]" + - name: Start headless display + uses: pyvista/setup-headless-display-action@v4 + - name: Verify external renderer tools + if: github.event_name != 'pull_request' + run: | + python - <<'PY' + from choreographer.browsers.chromium import Chromium + from postgkyl.render import _ffmpeg + + chrome = Chromium.find_browser(skip_local=False) + ffmpeg = _ffmpeg.resolve_ffmpeg() + if chrome is None: + raise SystemExit("Chrome/Chromium is required for Kaleido tests") + if ffmpeg is None: + raise SystemExit("ffmpeg is required for animation tests") + print(f"chrome: {chrome}") + print(f"ffmpeg: {ffmpeg}") + PY + - name: Run headless renderer tests + run: >- + python -m pytest + -m "render and not external_tool" + --strict-config --strict-markers + --timeout=120 + - name: Run external renderer tests + if: github.event_name != 'pull_request' + # These tests launch Chrome and ffmpeg. A separate invocation gives + # them a firm timeout and an unambiguous failure in the job log. + run: >- + python -m pytest + -m external_tool + --strict-config --strict-markers + --timeout=180 --timeout-method=signal + + package: + name: Build and install distribution artifacts + runs-on: ubuntu-latest + timeout-minutes: 60 steps: - - uses: actions/checkout@v3 - - uses: FedericoCarboni/setup-ffmpeg@v3 - id: setup-ffmpeg - with: - ffmpeg-version: release - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 - with: - python-version: ${{ matrix.python-version }} - - name: Install postgkyl - run: | - pip install -e .[adios,test] - - name: Test with pytest - run: | - pytest + - name: Check out source commit + uses: actions/checkout@v4 + with: + repository: ${{ env.SOURCE_REPOSITORY }} + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + - name: Verify source commit + shell: bash + run: test "$(git rev-parse HEAD)" = "$SOURCE_SHA" + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + - name: Install artifact tools and build prerequisites + run: >- + python -m pip install + "numpy==2.2.6" setuptools wheel + build==1.3.0 twine==6.2.0 + - name: Build sdist and wheel from the sdist + # `python -m build` creates the sdist first and then uses that sdist as + # the wheel source, testing that no untracked checkout files are needed. + run: python -m build --no-isolation + - name: Validate artifact metadata + run: python -m twine check --strict dist/* + - name: Remove access to build-time native libraries + shell: bash + run: | + # An absolute RPATH back into the checkout can make a broken wheel + # appear healthy. Hide the build tree before testing the installed + # artifact so the wheel must carry/find its own runtime libraries. + if [ -d "$GITHUB_WORKSPACE/gkeyll" ]; then + mv "$GITHUB_WORKSPACE/gkeyll" "$RUNNER_TEMP/build-only-gkeyll" + fi + - name: Smoke-test the installed wheel outside the checkout + run: scripts/smoke_wheel.sh dist/*.whl + - name: Upload distribution artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: distributions + path: dist/ + if-no-files-found: warn + retention-days: 14 diff --git a/.gitignore b/.gitignore index 7f886dfb..5fdab9bd 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,30 @@ dist/ postgkyl.egg-info/ src/postgkyl/version.py .DS_Store +tests/test_data/generated/ + +# vendored Gkeyll checkout, cloned at pip install time (scripts/build_gkeyll.sh) +/gkeyll/ + +# scratch files at the repo root only (unanchored patterns would silently +# ignore package sources under src/ and fixtures under tests/) +/*.gkyl +/*.md +!/AGENTS.md +/*.py +!/setup.py +/*.json +# built gpython/_gpython extension (scripts/build_gpython.sh) +src/postgkyl/gpython/_gpython.so +# bundled native library copied beside the extension +src/postgkyl/gpython/libg0core.so +# generated Gkeyll commit/build-date metadata (scripts/build_gpython.sh) +src/postgkyl/gpython/_build_info.py + +# scratch output from running examples/scripts/*.py directly +examples/scripts/output/ +.vscode/settings.json +.coverage +.coverage.* +coverage.xml +htmlcov/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..e8a04a35 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,27 @@ +minimum_pre_commit_version: "4.3.0" + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-merge-conflict + - id: check-toml + - id: check-yaml + - id: end-of-file-fixer + - id: trailing-whitespace + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.12.11 + hooks: + - id: ruff-check + + - repo: https://github.com/google/yapf + rev: v0.43.0 + hooks: + - id: yapf + + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v22.1.3 + hooks: + - id: clang-format + types_or: [c, c++] diff --git a/.style.yapf b/.style.yapf new file mode 100644 index 00000000..027356dc --- /dev/null +++ b/.style.yapf @@ -0,0 +1,5 @@ +[style] +based_on_style = pep8 +column_limit = 80 +indent_width = 2 +continuation_indent_width = 4 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..3340a41e --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,22 @@ +# Postgkyl + +You are a developer of the postgkyl postprocessing tool, which is used +for reading, analyzing and visualizing Gkeyll simulation data. +It gives scientists one composable API +while preserving the mathematical meaning. +You are helpful and aim to construct a maintainable package. +Push back if something is wrong. +You care most about readable code. You dislike repeating code multiple times, so refactoring and organization is a high priority. + +Keep changes locally understandable, with one owner for each fact and computation. +For code changes, read the design skill and the task-relevant skills below: + +- [Design](.agents/skills/postgkyl-design/SKILL.md): coding doctrine. +- [Architecture](.agents/skills/postgkyl-architecture/SKILL.md): layer ownership and imports. +- [Data](.agents/skills/postgkyl-data/SKILL.md): backends, representations, and conversions. +- [API and CLI](.agents/skills/postgkyl-api-cli/SKILL.md): public verbs and generated commands. +- [Native bridge](.agents/skills/postgkyl-native/SKILL.md): Gkeyll integration and ownership. +- [Development](.agents/skills/postgkyl-development/SKILL.md): setup, checks, and agent configuration. +- [Testing] (.agents/skills/postgkyl-testing/SKILL.md): How to write good unit tests. + +`.agents/` is shared through `.claude` and `.codex` symlinks; `CLAUDE.md` links here. diff --git a/AUTHORS.md b/AUTHORS.md index b12ab253..fc2760a7 100644 --- a/AUTHORS.md +++ b/AUTHORS.md @@ -3,17 +3,17 @@ ## Design and main code development - Ammar Hakim (PPPL) +- Maxwell Rosen (PPPL) (2.0 main contributions) - Petr Cagas (HZDR/CASUS) -## Contributors (ultra alphabetically) +## Contributors (ultra-alphabetically) -- Manaure Francisquez (PPPL) -- Luca Georgescu -- Jonathan Gorard (PPPL) - James Juno (PPPL) -- Noah Mandel (PPPL) -- Maxwell Rosen (PPPL) +- Jonathan Gorard (PPPL) - Liang Wang +- Luca Georgescu +- Manaure Francisquez (PPPL) +- Noah Mandel (PPPL) ## Acknowledgement diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 00000000..47dc3e3d --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000..dfbf7490 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,11 @@ +include scripts/build_gkeyll.sh +include scripts/build_gpython.sh +include scripts/gkeyll-revision +include scripts/smoke_wheel.sh +include src/postgkyl/gpython/csrc/_gpythonmodule.c + +# These are generated while building a wheel. Never let the order of local +# release commands leak a platform binary or stale provenance into an sdist. +exclude src/postgkyl/gpython/_build_info.py +exclude src/postgkyl/gpython/_gpython.so +exclude src/postgkyl/gpython/libg0core.so diff --git a/README.md b/README.md index 74d6b177..23908cf8 100644 --- a/README.md +++ b/README.md @@ -6,122 +6,296 @@ This is the Postgkyl project. It is both Python library and command-line tool designed to provide unified access to Gkeyll data together with a broad variety of analytical and visualization tools. -## Documentation +## Installation -Full documentation of the Gkeyll project is available at -[ReadTheDocs](http://gkeyll.rtfd.io). +Follow these steps to install the current source version of Postgkyl. Run +the commands one line at a time in a terminal, using the same terminal +throughout. These instructions use bash or zsh on Linux or macOS. On Windows, +first set up [Ubuntu in WSL](https://learn.microsoft.com/en-us/windows/wsl/install) +and use its terminal. -## Dependencies and Installation +You need an internet connection, Git (to download the code), and build tools +(to build part of Postgkyl). Install the tools for your system: -Postgkyl requires the following packages: +- **Ubuntu / Debian, including WSL:** run `sudo apt update`, then + `sudo apt install git build-essential`. +- **macOS:** run `xcode-select --install` and complete the installer. +- **Other Linux distributions:** install Git, Make, and a C compiler using + your distribution's package manager. -* [click](https://pypi.org/project/click/) -* [matplotlib](https://pypi.org/project/matplotlib/) -* [msgpack](https://pypi.org/project/msgpack/) -* [numpy](https://pypi.org/project/numpy/) -* [scipy](https://pypi.org/project/scipy/) -* [sympy](https://pypi.org/project/sympy/) -* [tables](https://pypi.org/project/tables/) +Download Postgkyl and enter its folder: -Note that Posgkyl currently does not work with NumPy >= 2.0; the update is in -the works. In addition, there are two optional dependencies: +```bash +git clone https://github.com/ammarhakim/postgkyl.git +cd postgkyl +``` -* [adios2](https://pypi.org/project/adios2/) -* [pytest](https://pypi.org/project/pytest/) +### 1. Create an environment -ADIOS 2 is required for reading Gkeyll 2 `bp` output files and it is not needed -when working only with `gkylzero`. [pytest](https://docs.pytest.org/en/stable/) -is required only for developers. +An environment keeps Postgkyl's Python packages separate from those used by +other projects. Choose **one** of the following options, then continue to +step 2. If you are new to Python environments, use **mamba**. -### Setting up virtual environment (recommended) +#### pyenv + +[pyenv](https://github.com/pyenv/pyenv#installation) lets you install a +specific Python version. If it is not installed, follow its installation +instructions, including the +[Python build prerequisites](https://github.com/pyenv/pyenv/wiki#suggested-build-environment) +and shell setup, then reopen your terminal and return to the `postgkyl` +folder. + +Install Python 3.12, select it for this folder, and use Python's built-in +`venv` tool to create and activate the environment: + +```bash +pyenv install 3.12 +pyenv local 3.12 +python -m venv .venv +source .venv/bin/activate +``` -We strongly recommend creating a virtual Python environment for everybody -working with more than one Python project (this includes even using both -Postgkyl and Sphinx). The two recommended options are -[venv](https://docs.python.org/3/library/venv.html) and -[mamba](https://mamba.readthedocs.io/en/latest/). +#### mamba -With `venv`, one can create the virtual environment with: +If you do not have mamba, install +[Miniforge](https://github.com/conda-forge/miniforge#install), which includes +it. Allow the installer to initialize your shell, then reopen your terminal +and return to the `postgkyl` folder. Create and activate the environment: ```bash -python -m venv /path/to/new/virtual/environments/pgkyl +mamba env create -f environment.yml +mamba activate pgkyl ``` -then activate it with: +If you already use **conda** (for example, through Anaconda or Miniconda), +you can use `conda` in place of `mamba` in both commands. + +### 2. Install dependencies + +Dependencies are the other packages Postgkyl needs. The two configuration +files have different jobs: -| bash/zsh | `source /bin/activate` | -| fish | `source /bin/activate.fish` | -| csh/tcsh | `source /bin/activate.csh` | +- [environment.yml](environment.yml) creates the mamba/conda environment + with Python, pip (the Python package installer), and setuptools (a build + tool). The pyenv/venv option above sets up Python and pip without this file. +- [pyproject.toml](pyproject.toml) lists the packages Postgkyl uses, such as + NumPy and Matplotlib, and the optional developer tools. pip reads this file + when installing Postgkyl, so these dependency lists are maintained here. -and deactivate with: +With your environment active, install NumPy and the Python build tools: ```bash -deactivate +python -m pip install --upgrade "numpy>=2.2.6" setuptools wheel ``` -With `mamba`, one can create the virtual environment with: +NumPy must be installed before building Postgkyl. The remaining dependencies +will be installed automatically in step 3; you do not need to install them +one by one. + +### 3. Install Postgkyl + +From the `postgkyl` folder, run: ```bash -mamba create -n pgkyl +python -m pip install --no-build-isolation . ``` -then activate with: +The final `.` means "install from this folder." Keep `--no-build-isolation` +so Postgkyl builds using the NumPy you just installed. This step also +downloads and builds the Gkeyll bridge automatically and may take several +minutes. + +Check the installation: ```bash -mamba activate pgkyl +pgkyl --version +pgkyl --help ``` -and deactivate with: +The first command prints version information; the second lists the available +commands. Each time you open a new terminal, activate your environment again: +run `source .venv/bin/activate` from the `postgkyl` folder for pyenv/venv, or +`mamba activate pgkyl` (or `conda activate pgkyl`) for mamba/conda. + +## Documentation + +Full documentation of the Gkeyll project, including Postgkyl, is available at +[ReadTheDocs](https://gkeyll.readthedocs.io/en/latest/postgkyl/index.html). The repository also contains +[examples](examples/README.md) and [notebooks](notebooks/README.md). + +The website generates this section from Postgkyl's `main` branch. Guides live +in [docs/source](docs/source), reference pages come from the Python API and +compiled CLI, and figures come from the tested example scripts. See the +[documentation build instructions](docs/source/contributing.rst) for a local +preview and [integration notes](docs/integration-plan.md) for the website setup. + +For help with a particular command, add `--help`, for example: ```bash -mamba deactivate +pgkyl interpolate --help ``` -Note that with `mamba`, one can also use the provided `environment.yml` file, -which also includes dependency specifications: +### Additional installation notes + +To install the published version from [PyPI](https://pypi.org/project/postgkyl/), +complete the environment and dependency steps above, then replace the install +command in step 3 with: ```bash -mamba env create -f environment.yml +python -m pip install --no-build-isolation postgkyl ``` -### Installing Postgkyl +To leave an environment, run `deactivate` for pyenv/venv, `mamba deactivate` +for mamba, or `conda deactivate` for conda. In shells other than bash/zsh, +venv activation uses `source .venv/bin/activate.fish` for fish or +`source .venv/bin/activate.csh` for csh/tcsh. + +Installing with pip does not require changes to `PYTHONPATH`. If you +previously added a Postgkyl checkout to that variable, remove that entry so +Python uses the installed package. + +#### Gkeyll bridge -The Postgkyl itself is installed with `pip`.[^1] Developers and uses who want to -have the most up-to-date version should install Postgkyl from the source code: +The Gkeyll bridge (`gpython`) connects Postgkyl to Gkeyll's compiled code for +native `.gkyl` reading, interpolation, integration, and DG arithmetic. +Installing from source builds it automatically. A prebuilt wheel includes +the bridge already. + +During a source build, `setup.py` runs `scripts/build_gkeyll.sh`, which: + +1. Downloads the [Gkeyll](https://github.com/ammarhakim/gkeyll) revision + recorded in `scripts/gkeyll-revision` into `gkeyll/`. +2. Builds its core library with the bundled LAPACK implementation. No + separate MPI, CUDA, SuperLU, Lua, or system LAPACK installation is needed. +3. Builds the Python extension and bundles the core library beside it, so + the installed package can run without the Gkeyll source folder. + +The build needs Git, Make, a C compiler, and network access. It uses `cc` by +default. To select another installed compiler, for example GCC, run: ```bash -git clone https://github.com/ammarhakim/postgkyl.git -cd postgkyl -pip install -e .[adios,test] +CC=gcc python -m pip install --no-build-isolation . ``` -Alternatively, Postgkyl can be installed directly from [PyPI](https://pypi.org/project/postgkyl/): +Always use `--no-build-isolation` when building from source so the bridge +builds against the NumPy in your active environment. A different NumPy at +build time can cause import errors or crashes. After changing NumPy, rebuild +the bridge in that environment. + +Check whether the bridge is available: ```bash -pip install -e postgkyl[adios,test] +python -c "from postgkyl import gpython; print(gpython.available())" ``` -Note that ADIOS2 is not available on PyPI for Mac OSX; therefore, Mac users who -want to use it need to install the dependency from elsewhere, for example, using -the above-mentioned `mamba` and then do *not* use the `adios` tag with `pip`. +This should print `True`. If it prints `False`, get the error details with: -#### Optional path to Gkeyll +```bash +python -c "from postgkyl import gpython; gpython.require()" +``` -Some features require telling postgkyl the location of the `core` library -of [Gkeyll](https://github.com/ammarhakim/gkeyll). After installing postgkyl you -can optionally write a postgkyl `config` file storing the location of Gkeyll's -`gkylsoft/` installation library by invoking the postgkyl config command, e.g. +A failed source build stops installation. An installation with an unavailable +bridge can still read files through the Python reader, but operations that +require the bridge raise an error. + +To rebuild and reinstall from the repository folder with your environment +active, run: + +```bash +python -m pip install --no-build-isolation . +``` + +For an editable developer installation (see below), you can rebuild in place +with `PYTHON=python scripts/build_gkeyll.sh`. If the Gkeyll core library is +already built and only the Python extension needs rebuilding, use +`PYTHON=python scripts/build_gpython.sh` instead. + +## Developing for Postgkyl + +Complete the installation steps above, then run this from the `postgkyl` +folder with your environment active: + +```bash +python -m pip install --no-build-isolation -e '.[test]' +``` + +The `-e` option makes the installation use your source files directly, so +Python edits take effect without reinstalling. The `[test]` option also +installs the testing, formatting, and packaging tools listed in +`pyproject.toml`. + +### pytest + +[pytest](https://docs.pytest.org/) runs the automated tests. From the +`postgkyl` folder, run: + +```bash +python -m pytest tests/ +``` + +Add `-v` to see a separate result for each test. + +The default suite treats unexpected warnings as errors and uses strict marker +and configuration validation. Useful CI-equivalent subsets are: + +```bash +POSTGKYL_SKIP_GKEYLL_BUILD=1 pytest -m compatibility +POSTGKYL_REQUIRE_GKEYLL=1 pytest -m native +pytest -m "render and not external_tool" +pytest -m external_tool # invokes Chrome and/or ffmpeg +pytest -m "not external_tool" --cov=postgkyl --cov-branch --cov-fail-under=93 ``` -pgkyl config -g /gkylsoft/ -c ~/.postgkyl/gkylsoft_path + +The external-tool lane has explicit timeouts in CI. Native lanes set +`POSTGKYL_REQUIRE_GKEYLL=1`, turning a missing bridge into a session failure +instead of allowing the native test inventory to skip silently. + +For pure-Python compatibility testing, skip the native build at installation +time with `POSTGKYL_SKIP_GKEYLL_BUILD=1`. Use this only when testing the +`compatibility` subset; normal installations build the bridge. + +### Formatting + +After installing the developer tools above, enable the checks that run +before a Git commit and run them over all tracked files: + +```bash +pre-commit install +pre-commit run --all-files ``` -## Testing +pre-commit installs the pinned YAPF, clang-format, Ruff, and repository +checks. YAPF reads `.style.yapf`; clang-format reads `.clang-format`; Ruff +reads `pyproject.toml`. The automated pull-request checks use these same +tools and report any formatting changes needed. -Postgkyl utilizes [pytest](https://docs.pytest.org/) for testing. The tests can -be called manually from the root Postgkyl directory simply by using: +### API and CLI documentation + +Public command documentation lives on the Python function that implements the +operation. The equivalent `GData` spelling is a class-body alias to that same +function, so editor hover help, `help(pg.interpolate)`, +`help(data.interpolate)`, and `pgkyl interpolate --help` cannot maintain +separate descriptions. +The installed distribution includes a `py.typed` marker so language servers +consume these inline signatures and aliases from a virtual environment too. + +Command docstrings use `Args:` entries in Google style. Every CLI-visible +parameter needs one entry; command compilation rejects missing, duplicate, or +unknown parameter documentation. `tests/test_documentation.py` additionally +checks the public Python surface, static fluent aliases, source/runtime +docstring identity, and deterministic CLI lowering. Run it directly with: + +```bash +pytest tests/test_documentation.py +``` + +### Checking a release package + +Build a wheel (an installable package) and test it in a clean environment: ```bash -pytest [-v] +python -m build --no-isolation +scripts/smoke_wheel.sh dist/*.whl ``` ## Authors @@ -131,7 +305,3 @@ The full list of authors can be found [here](AUTHORS.md). ## License Postgkyl is distributed under the MIT License. - -[^1]: This does *not* require any additional modifications of `PYTHONPATH`. If - Postgkyl was used previously through `PYTHONPATH`, we strongly recommend - removing the path to the Postgkyl repository from the variable. diff --git a/docs/_ext/postgkyl_docs.py b/docs/_ext/postgkyl_docs.py new file mode 100644 index 00000000..9dfbcc3a --- /dev/null +++ b/docs/_ext/postgkyl_docs.py @@ -0,0 +1,22 @@ +"""Copy standalone interactive figures beside their referring tutorial pages.""" + +from pathlib import Path +import shutil + + +def copy_interactive(app, exception): + if exception is not None or app.builder.format != "html": + return + relative = Path(app.config.postgkyl_doc_root) / "interactive" + source = Path(app.srcdir) / relative + if source.is_dir(): + destination = Path(app.outdir) / relative + destination.mkdir(parents=True, exist_ok=True) + for path in source.glob("*.html"): + shutil.copyfile(path, destination / path.name) + + +def setup(app): + app.add_config_value("postgkyl_doc_root", ".", "html") + app.connect("build-finished", copy_interactive) + return {"parallel_read_safe": True, "parallel_write_safe": True} diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000..5fd1bd37 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,15 @@ +"""Standalone preview of the same source subtree staged into gkyl-doc.""" + +from pathlib import Path +import sys + +sys.path.insert(0, str(Path(__file__).resolve().parent / "_ext")) +project = "Postgkyl" +extensions = [ + "sphinx.ext.autodoc", "sphinx.ext.napoleon", "myst_parser", "postgkyl_docs" +] +html_theme = "furo" +root_doc = "index" +exclude_patterns = ["_inputs/**", "_ext/**"] +autodoc_typehints = "none" +nitpicky = False diff --git a/docs/integration-plan.md b/docs/integration-plan.md new file mode 100644 index 00000000..7dbe3f5e --- /dev/null +++ b/docs/integration-plan.md @@ -0,0 +1,41 @@ +# Documentation integration + +The website builds documentation from **Postgkyl main**, as requested. There is +no pinned Postgkyl dependency or submodule. The code, function docstrings, +compiled CLI, tutorials, fixtures, and figures come from one checkout per +build. Local and pull-request previews use the proposed working checkout. + +- `docs/source/` owns guides and tutorial narratives. +- `scripts/build_docs.py` runs the examples and prepares a Sphinx source tree. +- `docs/conf.py` builds a standalone preview of that tree; `docs/_ext/` copies + interactive HTML beside its referring pages in either website layout. +- `examples/` owns executable Python and CLI tutorials; published code is + included directly from those files. +- `examples/figure_commands.json` owns the paired CLI pipelines. + `examples/compare_interfaces.py` executes scripts and CLI independently, + compares pixels, GIF frames/timings, and Plotly data/layout/configuration, + and produces the report used by the website. +- `tests/generate_test_data.py` owns synthetic fixtures, including the + shock-tube state, exponential energy history, travelling waves, and 3D Gaussian. +- `tests/test_docs_build.py` checks strict Sphinx builds, command coverage, + per-command API pages, navigation, portable downloads, and preservation of unrelated directories. +- `.github/workflows/docs.yml` runs the documentation checks on Python 3.12 + for pull requests and main pushes, and uploads an HTML preview. + +In gkyl-doc, `scripts/prepare_postgkyl.py` fetches main, installs that checkout, +and stages the generated section. `make html`, Read the Docs, and the host's +GitHub Actions use that script. The host also tests daily to catch upstream +changes without requiring a host commit. Its README describes the optional +Read the Docs token for publishing after those scheduled checks. + +Merge the Postgkyl side first, then the host integration. Hosted build time and +account configuration must be checked on Read the Docs; the local integration +uses an explicit checkout so both sides can be tested before merging. + +The hosting configuration selects Python 3.12. Builds reject Sphinx warnings, +missing native support, broken examples, and a mismatched installed checkout. +A recorded source SHA explains a particular generated result; each normal +website build still fetches main afresh. + +See `docs/source/contributing.rst` for local commands and the gkyl-doc README +for the complete website build. diff --git a/docs/source/animation.rst b/docs/source/animation.rst new file mode 100644 index 00000000..91726d32 --- /dev/null +++ b/docs/source/animation.rst @@ -0,0 +1,29 @@ +Load many files, collect, and animate +======================================= + +The generated travelling-wave examples represent the same positive density +wave at successive times. These analytic fixtures demonstrate file handling +and visualization; they are not numerical simulation results. Each file has +its own frame number and time metadata. + +Download the :download:`example bundle `. +The zero-padded names (``travelling_wave_000.gkyl`` through frame 015) make +lexicographic sorting match time order. The script uses ``sorted(Path.glob(...))``; +the CLI accepts the quoted wildcard. Use ``sort`` with your own files if +their names or input order do not match time order. + +``collect(frames)`` makes a single dataset whose leading grid axis is time. +A 1-D sequence therefore becomes a 2-D space–time plot. Its timestamps are +sorted using the files' metadata. Interpolate before collecting. + +``animate(frames)`` instead draws the original individual frames in sequence. +Do not pass the collected space–time dataset as though it were a sequence. +The example keeps a fixed value range across all frames, so a changing color +does not merely reflect a changing scale. The second animation shows a 2-D +travelling wave. ``saveframes`` also saves each frame as a PNG; ``fps`` controls +playback speed, independently of the physical times stored in the files. + +Both GIFs below are generated and checked frame by frame during the build. +Use the browser's image controls or open the images separately to inspect them. + +.. include:: _pairs/07_collect_animate.inc diff --git a/docs/source/arithmetic-representation.rst b/docs/source/arithmetic-representation.rst new file mode 100644 index 00000000..030b1899 --- /dev/null +++ b/docs/source/arithmetic-representation.rst @@ -0,0 +1,14 @@ +Choose the arithmetic representation +==================================== + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs below are both generated and compared when this website is built. + +The script compares weak DG algebra with pointwise operations after +interpolation. A projected product does not in general preserve all terms +of the full polynomial product. A nonzero divisor alone does not guarantee +an exact multiply/divide round trip. + +.. include:: _pairs/02_arithmetic_and_numpy.inc diff --git a/docs/source/cli-tutorial.rst b/docs/source/cli-tutorial.rst new file mode 100644 index 00000000..6d732aa4 --- /dev/null +++ b/docs/source/cli-tutorial.rst @@ -0,0 +1,24 @@ +Command-line workflows +======================== + +Every figure in :doc:`examples` has both a complete Python script +and a copyable ``pgkyl`` command. Their generated outputs appear together, +with the result of the build's equivalence check. See +:doc:`interface-equivalence` for the complete comparison report. + +The CLI keeps a working set of datasets: loading adds files, transformations +operate on them, ``collect`` combines frames, and ``plot`` or ``animate`` +renders the result. A quoted filename wildcard loads multiple matching files. +The Python spelling is a list of loaded datasets passed to the corresponding +function. The animation tutorial demonstrates both forms. + +Command and option names preserve Python underscores. A boolean flag such as +``--no_show`` means ``True``; explicit ``--no_show True`` and ``--no_show False`` +are also accepted. ``--saveas`` chooses the output path. The gallery's CLI +commands use ``output/``; the scripts use ``examples/scripts/output/`` by +default, or the directory specified by ``PGKYL_EXAMPLE_OUTPUT``. + +Use :doc:`reference/cli` to inspect every command and +:doc:`reference/api` for the corresponding Python calls. The example bundle +also contains the longer, executable CLI walkthrough in +``examples/cli_tutorial.md``. diff --git a/docs/source/concepts.rst b/docs/source/concepts.rst new file mode 100644 index 00000000..8c382264 --- /dev/null +++ b/docs/source/concepts.rst @@ -0,0 +1,64 @@ +Working with data +=================== + +Coordinates, components, and metadata +--------------------------------------- + +``pg.load(path)`` returns a ``GData`` containing coordinates, values, and +metadata in ``ctx``. Inspect the file before interpreting component numbers: +a five-moment fluid state and a gyrokinetic distribution do not use the same +layout. Diagnostics know the equation-specific interpretation. + +Basis type, polynomial order, and ``value_form`` are properties of the data. +They are read from the file or supplied once to ``load`` when missing. +Downstream operations use that metadata. Do not guess a higher-order basis +for a file whose header lacks this information. + +Coefficients and point values +------------------------------- + +``value_form`` distinguishes three representations: + +* ``modal`` stores DG expansion coefficients. Multiplication and division + use weak DG operations; coefficients are not samples for NumPy ufuncs. +* ``nodal`` stores values at basis nodes. +* ``quad`` stores values at Gauss–Legendre quadrature points. + +Nodal and quadrature values support pointwise arithmetic, including native +data. The native/NumPy backend and the value representation are separate +facts. ``to_modal()``, ``to_nodal()``, and ``to_quad()`` explicitly change +representation. ``interpolate()`` creates a new NumPy-backed field on a +refined mesh for selection, plotting, and general analysis. + +For nonlinear operations on a modal field, ``apply(fn, num_quad=...)`` spells +out evaluation at quadrature points followed by projection back to modal +coefficients. Mixing representations in arithmetic raises an error. + +Integration and cuts +---------------------- + +Integrating a modal field uses the DG representation directly. +``average`` and partial ``integrate`` can reduce dimensionality while keeping +a modal dataset for further operations. ``eval_at_coord_proj`` evaluates at +specified coordinates and projects into the basis of the surviving directions. + +After interpolation, ``select(comp=...)`` selects components and +``select(z0=..., z1=...)`` selects positions. A selected axis may retain a +singleton dimension. Rendering squeezes it as needed. +Use ``local_poly()`` when the plot must preserve jumps between DG cells. + +Composition +------------- + +Most transformations return a new dataset; ``inplace=True`` explicitly +requests mutation. ``plot`` returns a figure, and full integration returns +integrated values. A chain ends when it reaches a terminal result. + +``GDataGroup`` broadcasts operations over its members. CLI pipelines similarly +maintain a working set of loaded datasets; use ``evaluate`` to combine them +with an RPN expression. Public boolean parameters default to ``False``: +``--no_show`` means ``True`` and disables the plot window. Explicit +``--no_show True`` and ``--no_show False`` are also accepted. + +See :doc:`reference/api` for exact arguments and :doc:`reference/cli` for +the generated command spellings. diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst new file mode 100644 index 00000000..19748a87 --- /dev/null +++ b/docs/source/contributing.rst @@ -0,0 +1,58 @@ +Troubleshooting and documentation development +=============================================== + +If an operation reports missing Gkeyll support, run: + +.. code-block:: bash + + python -c "from postgkyl import gpython; gpython.require()" + +Follow the bridge rebuild instructions in :doc:`installation`. NumPy must +be installed before building the bridge, and the same NumPy ABI must be used +at runtime. If an operation instead rejects modal coefficients, use the +explicit representation change appropriate to the calculation; see +:doc:`concepts`. + +For a machine without a display, set ``MPLBACKEND=Agg`` and use +``no_show=True`` or ``--no_show``. Example scripts already save their plots +without opening windows. + +Build the documentation +------------------------- + +From a Postgkyl checkout, using Python 3.12: + +.. code-block:: bash + + python -m pip install --upgrade numpy setuptools wheel + python -m pip install --no-build-isolation -e '.[docs,test]' + POSTGKYL_REQUIRE_GKEYLL=1 MPLBACKEND=Agg python -m pytest tests/test_examples.py tests/test_documentation.py tests/test_docs_build.py + python scripts/build_docs.py + python -m sphinx -W --keep-going -b html -c docs build/docs/source build/docs/html + +Open ``build/docs/html/index.html``. The preparation step regenerates fixtures +and figures and replaces only an output directory bearing its ownership marker. +The example bundle preserves all paths used by the scripts and CLI walkthrough. + +``examples/figure_commands.json`` owns the CLI commands paired with the scripts. +``examples/compare_interfaces.py`` executes both interfaces and fails on changed +pixels, animation timing, or Plotly data/layout. The generated website includes +that comparison report and both outputs. To run just the paired gallery: + +.. code-block:: bash + + python examples/compare_interfaces.py + +The PyVista examples require OpenGL even when no window is shown. Headless +CI installs Mesa/EGL and sets ``VTK_DEFAULT_OPENGL_WINDOW=vtkEGLRenderWindow`` +and ``LIBGL_ALWAYS_SOFTWARE=1``. See :doc:`interactive` for local setup. + +Edit guides in ``docs/source/``, examples in ``examples/``, and API descriptions +in implementing function docstrings. The command reference reads the actual +compiled CLI. Both websites use this same preparation step; generated files +are not maintained by hand. + +The Gkeyll host fetches Postgkyl ``main`` on each build. Its scheduled GitHub +Actions check detects breakage from upstream changes even when the host +repository has no new commits. Postgkyl pull requests build their proposed +code and documentation together before merge. diff --git a/docs/source/examples.rst b/docs/source/examples.rst new file mode 100644 index 00000000..55eab39f --- /dev/null +++ b/docs/source/examples.rst @@ -0,0 +1,33 @@ +Examples +======== + +Each example has its own page with a complete Python script, the equivalent +``pgkyl`` commands, and their generated results. + +Install Postgkyl using :doc:`installation`, then download and unzip the +:download:`example bundle `. +Run the scripts and commands from the extracted directory so their data paths +resolve correctly. + +.. toctree:: + :maxdepth: 1 + + quickstart + profile-comparison + fluid-pressure + gyrokinetic-distribution + physical-rz + exponential-growth + arithmetic-representation + animation + plotly + pyvista + manual-arrays + +For help adapting the commands and comparing their results: + +.. toctree:: + :maxdepth: 1 + + cli-tutorial + interface-equivalence diff --git a/docs/source/exponential-growth.rst b/docs/source/exponential-growth.rst new file mode 100644 index 00000000..6fc0b3ea --- /dev/null +++ b/docs/source/exponential-growth.rst @@ -0,0 +1,16 @@ +Measure exponential growth +========================== + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs below are both generated and compared when this website is built. + +The generator supplies an analytic energy history with a known exponential +rate. Fit the logarithm of positive energy to time; its slope is the energy +growth rate. If energy is proportional to squared mode amplitude, the +amplitude growth rate is half this slope. This exact example checks the +calculation; real data require selecting a justified growth interval and +examining residuals before saturation dominates. The script checks the known rate and residuals numerically, then plots the reconstructed fitted energy. + +.. include:: _pairs/06_growth.inc diff --git a/docs/source/fluid-pressure.rst b/docs/source/fluid-pressure.rst new file mode 100644 index 00000000..7f211cd4 --- /dev/null +++ b/docs/source/fluid-pressure.rst @@ -0,0 +1,15 @@ +Recover pressure from conserved moments +======================================= + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs below are both generated and compared when this website is built. + +A fluid file stores density, momentum components, and total energy. +``five_moment.pressure`` subtracts kinetic energy and applies the specified +ratio of specific heats. The generated fixture is a stationary shock-tube +initial condition in dimensionless units, with gamma = 5/3. It does not show +the evolved shock, contact, or rarefaction. + +.. include:: _pairs/03_diagnostics_five_moment.inc diff --git a/docs/source/gyrokinetic-distribution.rst b/docs/source/gyrokinetic-distribution.rst new file mode 100644 index 00000000..ebe50415 --- /dev/null +++ b/docs/source/gyrokinetic-distribution.rst @@ -0,0 +1,21 @@ +Inspect a gyrokinetic distribution +================================== + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs below are both generated and compared when this website is built. + +This example uses committed TCV output: ion Hamiltonian moments from +``rt_gk_tcv_iwl_adapt_source_1x2v_p1`` and an electron distribution plus geometry +from ``rt_gk_tcv_iwl_1x2v_p1``. These are separate runs; do not combine their +moments into one conservation calculation. Frame 250 identifies saved output, +not a time in seconds. + +``load_quantity`` resolves named quantities from available files. +``load_distf`` reconstructs the distribution using its matching Jacobians. +Coordinates remain computational in this example; the optional velocity +mapping is not enabled. Inspect the returned grid before assigning units or +interpreting a velocity-space slice. + +.. include:: _pairs/04_gyrokinetics.inc diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 00000000..30cbea4a --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,18 @@ +Postgkyl +======== + +Postgkyl reads Gkeyll output, computes derived quantities, and creates plots +through a Python library and the chainable ``pgkyl`` command-line tool. +Start with a complete example, then use the reference to adapt it to your run. + +.. toctree:: + :maxdepth: 1 + + installation + examples + Python API + Command reference + concepts + reference/quantities + contributing + provenance diff --git a/docs/source/interactive.rst b/docs/source/interactive.rst new file mode 100644 index 00000000..4f1ed12b --- /dev/null +++ b/docs/source/interactive.rst @@ -0,0 +1,19 @@ +:orphan: + +Interactive surfaces and volume rendering +=========================================== + +Download the :download:`example bundle ` +and run from the extracted directory. The small synthetic data files are +generated by ``tests/generate_test_data.py`` and already included in the bundle. +All coordinates and scalar amplitudes in these examples are dimensionless. + +These examples now have individual pages under :doc:`examples`. + +.. _plotly-2-d-height-surfaces-and-3-d-isosurfaces: + +* :doc:`plotly` + +.. _pyvista-contours-and-translucent-volumes: + +* :doc:`pyvista` diff --git a/docs/source/main.rst b/docs/source/main.rst new file mode 100644 index 00000000..570fe81d --- /dev/null +++ b/docs/source/main.rst @@ -0,0 +1,8 @@ +:orphan: + +Postgkyl documentation +======================== + +The current Postgkyl documentation starts at :doc:`index`. +Use that page for installation, executable tutorials, and the generated +Python and command-line reference. diff --git a/docs/source/manual-arrays.rst b/docs/source/manual-arrays.rst new file mode 100644 index 00000000..89ecb3ca --- /dev/null +++ b/docs/source/manual-arrays.rst @@ -0,0 +1,28 @@ +Access grids and values directly +================================== + +Postgkyl does not require you to express every calculation as a pipeline. +After loading a dataset, ``data.grid`` exposes its coordinate arrays and +``data.values`` exposes the stored numbers. What those numbers mean depends +on the representation: a raw modal file holds coefficients, while an +interpolated dataset holds field values. ``np.asarray(data)`` also exposes +point values; it deliberately refuses modal coefficients. + +Use ``.copy()`` when you want independent writable arrays. Never assume that +editing a view will leave its dataset unchanged. The example below copies +both the grid and values, performs a NumPy calculation, and uses ``clone`` +and ``push`` to create a plotted result while retaining the field metadata. + +An axis with one more coordinate than its value-array dimension contains +cell edges; use adjacent-edge midpoints to associate one coordinate with each +cell value. If the coordinates already align one-to-one with the values, +use them directly. ``np.meshgrid(..., indexing="ij")`` preserves the data's +axis order. The final array dimension indexes components, not space. + +This example converts density to a normalized perturbation. Its CLI equivalent +uses ``evaluate`` for the same arithmetic. The rendered pixels are checked +against the manually manipulated Python arrays. + +Download the :download:`example bundle ` first. + +.. include:: _pairs/10_manual_arrays.inc diff --git a/docs/source/physical-rz.rst b/docs/source/physical-rz.rst new file mode 100644 index 00000000..e63388c2 --- /dev/null +++ b/docs/source/physical-rz.rst @@ -0,0 +1,14 @@ +Plot in physical R–Z coordinates +================================ + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs below are both generated and compared when this website is built. + +Computational coordinates are useful for analysis, but a poloidal view places +the field in its physical geometry. Keep the companion +``rt_gk_tcv_nt_iwl_3x2v_p1-geo_int_mapc2p.gkyl`` beside the electron density +file so the operation can resolve it by the simulation prefix. + +.. include:: _pairs/05_gk_rz.inc diff --git a/docs/source/plotly.rst b/docs/source/plotly.rst new file mode 100644 index 00000000..69f793e3 --- /dev/null +++ b/docs/source/plotly.rst @@ -0,0 +1,21 @@ +Plotly: 2-D height surfaces and 3-D isosurfaces +=============================================== + +Download the :download:`example bundle ` +and run from the extracted directory. The small synthetic data files are +generated by ``tests/generate_test_data.py`` and already included in the bundle. +All coordinates and scalar amplitudes in these examples are dimensionless. + +For a two-dimensional field, ``plotly`` draws the scalar value as a height +above the x–y plane. A three-dimensional field instead uses nested isosurfaces +inside the volume. ``surface_count`` and ``opacity`` determine how many +surfaces appear and how well the interior can be seen. + +The HTML figures below are interactive: drag to rotate, scroll to zoom, and +hover to inspect values. Python and CLI outputs have exactly matching trace +data, layouts, and configurations; browser rasterization can depend on the +reader's graphics hardware. HTML IDs are randomly generated and are excluded +from that comparison. No Chrome/Kaleido installation is needed to generate +these HTML examples. Viewing them requires network access to the Plotly runtime CDN. + +.. include:: _pairs/08_plotly.inc diff --git a/docs/source/profile-comparison.rst b/docs/source/profile-comparison.rst new file mode 100644 index 00000000..95caa058 --- /dev/null +++ b/docs/source/profile-comparison.rst @@ -0,0 +1,14 @@ +Compare profiles across calculations +==================================== + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs below are both generated and compared when this website is built. + +Compare density, parallel velocity, and two temperatures while retaining +visibility into low-density tails. These are analytic synthetic profiles +generated for plotting tests, not evidence of simulation convergence. +The logarithmic side masks nonpositive values, including negative velocity. + +.. include:: _pairs/mirror_comparison.inc diff --git a/docs/source/pyvista.rst b/docs/source/pyvista.rst new file mode 100644 index 00000000..443a8f86 --- /dev/null +++ b/docs/source/pyvista.rst @@ -0,0 +1,23 @@ +PyVista: contours and translucent volumes +========================================= + +Download the :download:`example bundle ` +and run from the extracted directory. The small synthetic data files are +generated by ``tests/generate_test_data.py`` and already included in the bundle. +All coordinates and scalar amplitudes in these examples are dimensionless. + +``pyvista`` renders the same three-dimensional Gaussian with VTK. Contours +highlight selected density levels; volume rendering shows the continuous +interior. The examples save screenshots with ``no_show=True`` and disable +camera rotation with ``no_spin=True``. Omit ``no_show`` to open a local +interactive window. The field occupies ``[-2, 2]`` in each coordinate; +``hide_axes=True`` keeps these views focused on the density structure. The +volume example uses a linear opacity ramp to reveal the dense core. + +An OpenGL context is required even for off-screen screenshots. The documentation +CI uses Mesa/EGL software rendering. On headless Linux install ``libegl1`` and +``libgl1-mesa-dri`` and select ``VTK_DEFAULT_OPENGL_WINDOW=vtkEGLRenderWindow``. +Failures are reported; a missing graphics backend does not silently remove +these examples from the website. + +.. include:: _pairs/09_pyvista.inc diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst new file mode 100644 index 00000000..b71e472d --- /dev/null +++ b/docs/source/quickstart.rst @@ -0,0 +1,24 @@ +Inspect a field and extract a lineout +======================================= + +The first question after a simulation finishes is often: what is stored in +this file, and how does one component vary along a particular cut? + +Install Postgkyl using :doc:`installation`. Download and unzip the +:download:`example bundle ` into a working +directory. It contains the scripts and data with the same paths as the repository. +From that directory, run: + +.. code-block:: bash + + python examples/scripts/01_quickstart.py + +The script loads an analytic two-component coordinate map, interpolates its +DG coefficients, and plots the first component. This small synthetic field +makes it easy to see which direction and component a selection acts on. + +``comp=0`` selects a component; ``z1=0.5`` extracts a spatial cut. +The script also saves and reloads the selected dataset and checks its values. +The paired figures below are compared pixel by pixel during the build. + +.. include:: _pairs/01_quickstart.inc diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst new file mode 100644 index 00000000..715f3f91 --- /dev/null +++ b/docs/source/tutorials.rst @@ -0,0 +1,35 @@ +:orphan: + +Scientific workflows +====================== + +All scripts below run against the :download:`example bundle +`. Unzip it, activate the environment where +Postgkyl is installed, and run the commands from the extracted directory. +The Python and CLI outputs are generated and compared when this website is built. + +These examples now have individual pages under :doc:`examples`. + +.. _compare-profiles-across-calculations: + +* :doc:`profile-comparison` + +.. _recover-pressure-from-conserved-moments: + +* :doc:`fluid-pressure` + +.. _inspect-a-gyrokinetic-distribution: + +* :doc:`gyrokinetic-distribution` + +.. _plot-in-physical-rz-coordinates: + +* :doc:`physical-rz` + +.. _measure-exponential-growth: + +* :doc:`exponential-growth` + +.. _choose-the-arithmetic-representation: + +* :doc:`arithmetic-representation` diff --git a/environment.yml b/environment.yml index f7a1a8b7..60a526c4 100644 --- a/environment.yml +++ b/environment.yml @@ -1,15 +1,10 @@ -name: pgkylSrcH5 +name: pgkyl channels: - defaults - conda-forge dependencies: - - click>=8.1.7 - - matplotlib>=3.7.0 - - msgpack-python>=1.0.3 - - numpy>=1.24.4,<2 - - pytables>=3.8.0 - - pytest>=7.4.0 - - python>=3.11 - - scipy>=1.10.1 - - sympy>=1.12 - - h5py + - python>=3.10 + # Runtime and test dependency versions have one authoritative home in + # pyproject.toml. Install the project after activating this environment. + - pip + - setuptools>=61.0 diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 00000000..ccd49e95 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,80 @@ +# Postgkyl examples + +Two parallel tutorials over the same golden path -- `load -> interpolate -> +select -> plot` -- one per interface: + +- **`scripts/`** -- the fluent Python API (`import postgkyl as pg`). +- **`cli_tutorial.md`** -- the equivalent `pgkyl` command-line chains. + +Both are executable, and both are tested: `tests/test_examples.py` runs +every script and replays every command quoted in the markdown tutorial, so +an example that stops working (an API rename, a removed option) fails the +test suite instead of quietly rotting. Treat that file as the single source +of truth for "does the tutorial still work," not just this README. + +## Scripts + +| Script | What it covers | +| --- | --- | +| [`01_quickstart.py`](scripts/01_quickstart.py) | `pg.load` → `.interpolate()` → `.select()` → `.plot()` → `.save()`/reload | +| [`02_arithmetic_and_numpy.py`](scripts/02_arithmetic_and_numpy.py) | Weak DG algebra on raw modal data (`*`, `/`, `+`, `.integrate()`) vs. plain NumPy math after `.interpolate()`, and the guardrail between them | +| [`03_diagnostics_five_moment.py`](scripts/03_diagnostics_five_moment.py) | The `diagnostics` layer: equation-specific physics (`postgkyl.diagnostics.mom.five_moment`) on top of a `GData`, on a generated shock-tube initial state | +| [`04_gyrokinetics.py`](scripts/04_gyrokinetics.py) | The gyrokinetic diagnostics: `pg.gk.load_quantity` (named moments/geometry, resolved by naming convention) and `pg.gk.load_distf` (full distribution function), on the `rt_gk_tcv_iwl*` fixtures | +| [`05_gk_rz.py`](scripts/05_gk_rz.py) | The gyrokinetic R-Z operation: one-line fluent and functional calls plus projection reuse over multiple toroidal angles | +| [`06_growth.py`](scripts/06_growth.py) | Recover a known energy growth rate and inspect log-space residuals | +| [`07_collect_animate.py`](scripts/07_collect_animate.py) | Load many frames, collect a space–time diagram, and animate 1D/2D travelling waves | +| [`08_plotly.py`](scripts/08_plotly.py) | Interactive 2D height surfaces and 3D volume isosurfaces | +| [`09_pyvista.py`](scripts/09_pyvista.py) | Off-screen VTK isosurface and volume screenshots | +| [`10_manual_arrays.py`](scripts/10_manual_arrays.py) | Copy grids/values, manipulate NumPy arrays, and wrap the result | +| [`mirror_comparison.py`](scripts/mirror_comparison.py) | A four-panel algorithm-sensitivity figure from two analytic, symmetric 1-D p1 modal-serendipity datasets, with joined linear/log axes | + +Run one directly: + +```bash +pip install --no-build-isolation -e '.[test]' +MPLBACKEND=Agg PYTHONPATH=src python examples/scripts/01_quickstart.py +``` + +Each script prints what it's doing as it goes and asserts the invariants it +demonstrates, so a successful run ending in `... OK` is itself a (manual) +confirmation that the example still holds. Output files (PNGs, a `.gkyl` +round trip) land in `examples/scripts/output/` by default, or wherever +`PGKYL_EXAMPLE_OUTPUT` points if that environment variable is set (this is +how `tests/test_examples.py` redirects them into a temp directory instead of +the repo). + +## CLI + +See [`cli_tutorial.md`](cli_tutorial.md) -- inspecting a file, the +`interpolate`/`select`/`plot` chain, discontinuity-preserving plots with +`local_poly`, DynVector `info`/`fit`, the gyrokinetic loaders +(`gk_load_quantity`, `gk_load_distf`), the `gk_rz` +transformation, `save`, and the generated command inventory. + +## Running the tests + +```bash +pytest tests/test_examples.py -v +``` + +Both halves need a compiled Gkeyll (`libg0core.so`) to run -- every fixture +here is a native `.gkyl` file, so the tests are skipped (not failed) +when `postgkyl.gpython.available()` is `False`, matching the rest of the +test suite's `needs_gkeyll` convention. + +## Script/CLI figure equivalence + +`figure_commands.json` holds a `pgkyl` pipeline for every published figure. +`compare_interfaces.py` runs those pipelines and the Python scripts separately, +then compares decoded pixels and GIF timing, or Plotly traces/layout/configuration. +Both outputs and the passing comparison report appear on the website. + +```bash +python tests/generate_test_data.py +python examples/compare_interfaces.py +``` + +Results land in `build/figure-comparison/`, with CLI versions under `cli/`. +The PyVista examples need OpenGL; headless CI uses Mesa/EGL with +`VTK_DEFAULT_OPENGL_WINDOW=vtkEGLRenderWindow` and `LIBGL_ALWAYS_SOFTWARE=1`. +The Plotly HTML examples require no browser executable to generate them. diff --git a/examples/cli_tutorial.md b/examples/cli_tutorial.md new file mode 100644 index 00000000..b9551690 --- /dev/null +++ b/examples/cli_tutorial.md @@ -0,0 +1,172 @@ +# `pgkyl` CLI tutorial + +Every command below is real: it runs against fixture files under +`tests/test_data/` -- most committed directly, a few under +`tests/test_data/generated/` and synthesized by `tests/generate_test_data.py` +(run once, or on the first `pytest` invocation, which does it automatically) +-- and `tests/test_examples.py` replays each one (via `click.testing.CliRunner`, +from the repository root) as a regression check. If the CLI's surface ever +changes in a way that breaks one of these commands, that test fails -- this +file cannot silently drift out of date the way a hand-maintained tutorial can. + +Run any line yourself from the repository root, after `pip install --no-build-isolation -e '.[test]'` +and (for the `tests/test_data/generated/` fixtures) `python +tests/generate_test_data.py`. + +## 1. Inspect a file + +`info` is the "what is this?" command -- dimensions, components, grid, +value range, and the DG basis/order it was written with. + +```bash +pgkyl tests/test_data/rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl info +``` + +`print` displays stored values with 16-digit precision and singleton +dimensions squeezed. For DG files, these are modal coefficients; run +`interpolate` first to print field values. Use `--grid` (or `-g`) to print +each grid axis, or `--use TAG` to print only datasets with that tag. +Printing leaves the datasets available for the rest of the chain. + +```bash +pgkyl tests/test_data/generated/energy_dynvec.gkyl print +pgkyl tests/test_data/generated/distf_p2_0.gkyl interpolate print +pgkyl tests/test_data/generated/distf_p2_0.gkyl print --grid +pgkyl tests/test_data/generated/energy_dynvec.gkyl --tag energy print --use energy +``` + +## 2. The chain: interpolate -> select -> plot + +Raw `.gkyl` files hold DG *coefficients*; `interpolate` bridges them onto a +uniform mesh of plain values, `select` narrows down to one component (or one +coordinate slice), and `plot` renders it. Boolean options default to `False` +and imply `True` when supplied without a value; an explicit `True` or `False` +is also accepted. The CLI mirrors the Python API exactly. A default-on +behavior therefore has a negative, default-`False` API parameter and CLI +name, so `no_show=False` / `--no_show False` displays a Matplotlib plot while +`--no_show` (or `--no_show True`) runs headless. `--saveas` writes a +PNG. + +```bash +pgkyl tests/test_data/rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl \ + interpolate select --comp 0 plot --no_show --saveas out.png +``` + +Any unambiguous command-name prefix is accepted as a spelling-only alias +(`interp` -> `interpolate`, `sel` -> `select`): + +```bash +pgkyl tests/test_data/rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl \ + interp sel --comp 0 info +``` + +## 3. Discontinuity-preserving plots with `local_poly` + +`interpolate` produces a continuous refined mesh; `local_poly` instead +evaluates the DG polynomial cell-by-cell and splices a NaN at every +inter-cell interface, so a plot shows genuine discontinuities instead of +smoothing over them -- useful for shocks or anything with jumps at cell +boundaries. + +```bash +pgkyl tests/test_data/generated/distf_p2_0.gkyl \ + local_poly select --z1 0.0 --z2 0.0 plot --no_show --saveas out.png +``` + +## 4. DynVector inspection and `fit` + +`.gkyl` files without a spatial grid (diagnostics like a field-energy history) +are DynVectors: `info` summarizes the data, and `fit` fits a model to it (here, +a straight line to the series vs. time -- a basic fit, not an exponential growth-rate measurement). + +```bash +pgkyl tests/test_data/generated/energy_dynvec.gkyl info +pgkyl tests/test_data/generated/energy_dynvec.gkyl fit --fit_type linear +``` + +## 5. Combining datasets: `evaluate` + +`evaluate` runs a Reverse Polish Notation expression over every dataset +currently loaded: `fN` refers to the `N`-th one in load order (`f` alone +means `f0`), so `"f0 f1 -"` subtracts the second dataset from the first -- +e.g. two frames of the same field, to see how it changed between them. +(`fN[c]` selects component `c` of dataset `N`; don't confuse that with +indexing datasets themselves -- there is no `f[N]` form.) Data must be +`interpolate`d first, same as `select`/`plot`. + +```bash +pgkyl tests/test_data/generated/distf_p2_0.gkyl tests/test_data/generated/distf_p2_1.gkyl \ + interpolate evaluate "f0 f1 -" info +``` + +## 6. Gyrokinetics: pre-named quantities and distribution functions + +`gk_load_quantity` loads one of a registry of named gyrokinetic +quantities (listed in its generated `--quantity` choices) straight from a simulation's naming convention -- +no manual file paths. `--name` is the simulation's *name prefix* (not a path); +`--path` is the directory to look in. + +```bash +pgkyl gk_load_quantity --help +pgkyl gk_load_quantity --quantity geo_int_jacobtot_inv --species "" \ + --name rt_gk_tcv_iwl_1x2v_p1 --path tests/test_data info +``` + +`gk_load_distf` reconstructs a full distribution function from the saved +`Jf`-times-Jacobian(s) files (here `--name` *does* include the directory, since +the simulation name itself includes the directory): + +```bash +pgkyl gk_load_distf --name tests/test_data/rt_gk_tcv_iwl_1x2v_p1 \ + --species elc --frame 250 \ + --jacobtot_inv_file tests/test_data/rt_gk_tcv_iwl_1x2v_p1-geo_int_jacobtot_inv.gkyl \ + info +``` + +## 7. Map a gyrokinetic field to R-Z + +`gk_rz` is a data transformation: it interpolates one raw DG component and +maps it onto the physical poloidal plane. Geometry is inferred from the +field's filename, preferring nodal geometry and falling back to modal +`mapc2p` geometry. The CLI and Python calls below use the same operation and +defaults: + +```bash +pgkyl tests/test_data/rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl \ + gk_rz --nz_interp 2 info +``` + +```python +import postgkyl as pg + +mapped = pg.load( + "tests/test_data/rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl" +).gk_rz(nz_interp=2) +``` + +## 8. Saving to another format + +`save` writes the current dataset(s) out as `gkyl`/`txt`/`npy`/`vtk`. + +```bash +pgkyl tests/test_data/generated/distf_p2_0.gkyl save --out_name distf --extension npy +``` + +## 9. One API-derived command inventory + +Every loaded file becomes a dataset in the current chain. The command list is +compiled from the script API, so `--help` is the authoritative inventory and +every Python underscore remains an underscore in the CLI. + +```bash +pgkyl --help +``` + +## See also + +- `pgkyl --help` lists every registered command, grouped by section + (Verbs / Diagnostics / Render / Utility). +- `pgkyl --help` documents that command's options -- most carry a + worked example in their docstring, e.g. `pgkyl local_poly --help`. +- `examples/scripts/` is the Python-script equivalent of this tutorial (the + fluent `GData` API instead of the chained CLI). diff --git a/examples/compare_interfaces.py b/examples/compare_interfaces.py new file mode 100644 index 00000000..a10ae287 --- /dev/null +++ b/examples/compare_interfaces.py @@ -0,0 +1,100 @@ +"""Run the tutorial scripts and CLI pipelines and compare their actual outputs. + +PNG/GIF comparisons use decoded pixels (and GIF timings). Plotly comparisons +use the saved trace data, layout, and configuration, ignoring random HTML IDs. +No cross-machine image snapshots or renderer mocks are used. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import re +import shlex +import subprocess +import sys + +import numpy as np +from PIL import Image, ImageSequence + + +def plotly_spec(path: Path) -> list: + html = path.read_text() + matches = list(re.finditer(r'Plotly\.newPlot\(\s*"[a-f0-9-]+"\s*,', html)) + if len(matches) != 1: + raise ValueError(f"Expected one saved Plotly figure in {path}") + remaining = html[matches[0].end():] + result = [] + for _ in range(3): + remaining = remaining.lstrip(" \n\r\t,") + value, end = json.JSONDecoder().raw_decode(remaining) + result.append(value) + remaining = remaining[end:] + return result + + +def compare_outputs(python_path: Path, cli_path: Path) -> str: + if python_path.suffix == ".html": + if plotly_spec(python_path) != plotly_spec(cli_path): + raise AssertionError( + f"Plotly data/layout/configuration differ: {python_path.name}") + return "Identical Plotly traces, layout, and configuration" + with Image.open(python_path) as first, Image.open(cli_path) as second: + if first.n_frames != second.n_frames: + raise AssertionError(f"Frame counts differ: {python_path.name}") + for index, (left, right) in enumerate( + zip(ImageSequence.Iterator(first), ImageSequence.Iterator(second))): + np.testing.assert_array_equal( + np.asarray(left.convert("RGBA")), + np.asarray(right.convert("RGBA")), + err_msg=f"{python_path.name}: frame {index}") + if left.info.get("duration") != right.info.get("duration"): + raise AssertionError(f"Frame timings differ: {python_path.name}") + return f"Identical pixels ({first.n_frames} frame(s)) and timing" + + +def run_gallery(root: Path, output: Path) -> dict[str, str]: + output.mkdir(parents=True, exist_ok=True) + cli_output = output / "cli" + cli_output.mkdir(exist_ok=True) + commands = json.loads((root / "examples/figure_commands.json").read_text()) + env = {**os.environ, "MPLBACKEND": "Agg", "PGKYL_EXAMPLE_OUTPUT": str(output)} + scripts = sorted((root / "examples/scripts").glob("*.py")) + results = {} + for script in scripts: + if script.name.startswith("_"): + continue + subprocess.run([sys.executable, str(script)], cwd=root, env=env, check=True) + for name, command in commands.get(script.name, {}).items(): + arguments = shlex.split(command) + if arguments.pop(0) != "pgkyl": + raise ValueError(f"Expected a pgkyl command for {name}") + arguments = [ + arg.replace("{output}", str(cli_output)) for arg in arguments + ] + subprocess.run([ + sys.executable, "-c", "from postgkyl.cli.app import cli; cli()", + *arguments + ], + cwd=root, + env=env, + check=True) + results[name] = compare_outputs(output / name, cli_output / name) + print(f"{name}: {results[name]}", flush=True) + expected = {name for group in commands.values() for name in group} + if set(results) != expected: + raise AssertionError(f"Unexecuted figure pairs: {expected - set(results)}") + (output / "interface-comparison.json" + ).write_text(json.dumps(results, indent=2) + "\n") + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", + type=Path, + default=Path("build/figure-comparison")) + args = parser.parse_args() + run_gallery(Path(__file__).resolve().parents[1], args.output.resolve()) diff --git a/examples/figure_commands.json b/examples/figure_commands.json new file mode 100644 index 00000000..b8be100b --- /dev/null +++ b/examples/figure_commands.json @@ -0,0 +1,43 @@ +{ + "01_quickstart.py": { + "01_quickstart_comp0.png": "pgkyl tests/test_data/generated/2d_c2p_rot45_ms_p1.gkyl interpolate select --comp 0 plot --title x-component --no_show --dpi 100 --saveas '{output}/01_quickstart_comp0.png'", + "01_quickstart_lineout.png": "pgkyl tests/test_data/generated/2d_c2p_rot45_ms_p1.gkyl interpolate select --z1 0.5 plot --title 'lineout at y=0.5' --no_show --dpi 100 --saveas '{output}/01_quickstart_lineout.png'" + }, + "03_diagnostics_five_moment.py": { + "03_diagnostics_density.png": "pgkyl tests/test_data/generated/shock_tube_1d_p0.gkyl interpolate five_moment_density plot --title 'density (Sod shock tube)' --no_show --dpi 100 --saveas '{output}/03_diagnostics_density.png'", + "03_diagnostics_pressure.png": "pgkyl tests/test_data/generated/shock_tube_1d_p0.gkyl interpolate five_moment_pressure --gas_gamma 1.6666666666666667 plot --title 'pressure (Sod shock tube)' --no_show --dpi 100 --saveas '{output}/03_diagnostics_pressure.png'" + }, + "04_gyrokinetics.py": { + "04_gyrokinetics_M0.png": "pgkyl gk_load_quantity --quantity M0 --species ion --name rt_gk_tcv_iwl_adapt_source_1x2v_p1 --frame 250 --path tests/test_data/ plot --title 'Ion density from Hamiltonian moments' --xlabel 'Field-aligned computational coordinate' --ylabel '$M_{0i}$ (m$^{-3}$)' --no_show --dpi 100 --saveas '{output}/04_gyrokinetics_M0.png'", + "04_gyrokinetics_jacobtot_inv.png": "pgkyl gk_load_quantity --quantity geo_int_jacobtot_inv --species '' --name rt_gk_tcv_iwl_1x2v_p1 --path tests/test_data/ plot --title '$(J B)^{-1}$' --no_show --dpi 100 --saveas '{output}/04_gyrokinetics_jacobtot_inv.png'", + "04_gyrokinetics_distf_slice.png": "pgkyl gk_load_distf --name tests/test_data/rt_gk_tcv_iwl_1x2v_p1 --species elc --frame 250 --jacobtot_inv_file tests/test_data/rt_gk_tcv_iwl_1x2v_p1-geo_int_jacobtot_inv.gkyl select --z2 0.0 plot --title 'Electron distribution: lowest mu slice' --xlabel 'Field-aligned computational coordinate' --ylabel 'Computational parallel velocity' --clabel 'Distribution [simulation units]' --no_show --dpi 100 --saveas '{output}/04_gyrokinetics_distf_slice.png'" + }, + "05_gk_rz.py": { + "05_gk_rz.png": "pgkyl tests/test_data/rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl gk_rz --z_axis 0.0 --phi_tor 0.0 --nz_interp 2 plot --title 'Electron density in the poloidal plane' --xlabel 'R [m]' --ylabel 'Z [m]' --clabel '$n_e$ [m$^{-3}$]' --fixaspect --no_show --dpi 100 --saveas '{output}/05_gk_rz.png'" + }, + "06_growth.py": { + "06_growth.png": "pgkyl tests/test_data/generated/exponential_energy.gkyl evaluate 'f log' fit --fit_type linear evaluate 'f1 exp' plot --title 'Fitted energy growth' --xlabel 'Time [normalized]' --ylabel 'Energy [normalized]' --logy --no_show --dpi 100 --saveas '{output}/06_growth.png'" + }, + "mirror_comparison.py": { + "mirror_comparison.png": "pgkyl tests/test_data/generated/mirror_comparison_2em4_1d_ms_p1.gkyl tests/test_data/generated/mirror_comparison_2em5_1d_ms_p1.gkyl interpolate plot --figure 0 --legend_labels 2e-4 --legend_labels 2e-5 --legend_subplot 0 --legend_loc best --title 'Comparing $\\alpha$ = 2e-4 to $\\alpha$ = 2e-5' --xlabel 'z [m]' --figsize 12 10 --split_linear_log --split_point 0 --split_log_side right --split_width_ratios 1 1 --split_gap 0 --split_legend_side log --split_log_nonpositive mask --subplot_ylabels 'Density [$\\mathrm{m}^{-3}$],$U_\\parallel$ [m/s],$T_\\parallel$ [keV],$T_\\perp$ [keV]' --no_show --dpi 150 --saveas '{output}/mirror_comparison.png'" + }, + "07_collect_animate.py": { + "07_collect.png": "pgkyl 'tests/test_data/generated/travelling_wave_*.gkyl' interpolate collect plot --title 'Travelling wave: space\u2013time diagram' --xlabel Time --ylabel x --clabel 'Density [normalized]' --no_show --dpi 100 --saveas '{output}/07_collect.png'", + "07_wave.gif": "pgkyl 'tests/test_data/generated/travelling_wave_*.gkyl' interpolate animate --saveframes '{output}/07_wave_frame' --fps 8 --dpi 90 --figsize 6 4 --no_show --saveas '{output}/07_wave.gif'", + "07_surface.gif": "pgkyl 'tests/test_data/generated/wave_surface_*.gkyl' interpolate animate --saveframes '{output}/07_surface_frame' --fps 6 --dpi 90 --figsize 6 4 --no_show --saveas '{output}/07_surface.gif'" + }, + "08_plotly.py": { + "08_surface.html": "pgkyl tests/test_data/generated/wave_surface_000.gkyl interpolate plotly --background light --title 'Travelling-wave height surface' --xlabel x --ylabel y --zlabel 'Density [normalized]' --clabel Density --saveas '{output}/08_surface.html'", + "08_volume.html": "pgkyl tests/test_data/generated/gaussian_volume.gkyl interpolate plotly --background light --title 'Gaussian density isosurfaces' --xlabel x --ylabel y --zlabel z --clabel Density --surface_count 8 --opacity 0.15 --cmin 0.1 --cmax 0.9 --saveas '{output}/08_volume.html'" + }, + "09_pyvista.py": { + "09_isosurfaces.png": "pgkyl tests/test_data/generated/gaussian_volume.gkyl interpolate pyvista --contour_levels 5 --no_show --no_spin --title 'Gaussian density isosurfaces' --xlabel x --ylabel y --zlabel z --clabel Density --theme document --saveas '{output}/09_isosurfaces.png' --camera_elevation 10 --camera_azimuth 20 --hide_axes", + "09_volume.png": "pgkyl tests/test_data/generated/gaussian_volume.gkyl interpolate pyvista --volume --no_show --no_spin --title 'Gaussian density volume' --xlabel x --ylabel y --zlabel z --clabel Density --theme document --saveas '{output}/09_volume.png' --camera_elevation 10 --camera_azimuth 20 --hide_axes --opacity linear" + }, + "10_manual_arrays.py": { + "10_manual.png": "pgkyl tests/test_data/generated/wave_surface_000.gkyl interpolate evaluate 'f 1 - 0.6 /' plot --title 'Normalized density perturbation' --xlabel x --ylabel y --clabel '(n - 1) / 0.6' --no_show --dpi 100 --saveas '{output}/10_manual.png'" + }, + "02_arithmetic_and_numpy.py": { + "02_arithmetic.png": "pgkyl tests/test_data/generated/2d_c2p_stretch_ms_p1.gkyl interpolate evaluate 'f sq 2 * sqrt' plot --title 'Componentwise magnitude' --no_show --dpi 100 --saveas '{output}/02_arithmetic.png'" + } +} diff --git a/examples/scripts/01_quickstart.py b/examples/scripts/01_quickstart.py new file mode 100644 index 00000000..4d6b96e9 --- /dev/null +++ b/examples/scripts/01_quickstart.py @@ -0,0 +1,70 @@ +"""Quickstart: the golden script path. + + pg.load(...).interpolate().select(...).plot() + +This walks through the same chain the CLI uses (see +``examples/cli_tutorial.md``), one step at a time: + +1. ``pg.load`` reads a ``.gkyl`` file. Raw DG data lands in the *modal* + backend -- coefficients, not values you can plot or feed to NumPy. +2. ``.interpolate()`` is the one-way bridge to a uniform mesh of plain + values (the *numpy* backend). +3. ``.select()`` picks out one component/coordinate. +4. ``.plot()`` renders it (a terminal verb -- returns the Matplotlib + ``Figure``, which this script then saves to disk). +5. ``.save()`` writes a dataset back out; reloading it recovers the same + values, byte for byte. + +Run directly: + MPLBACKEND=Agg PYTHONPATH=src python examples/scripts/01_quickstart.py +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") # headless-safe; drop this line to see the plot windows + +import postgkyl as pg +import numpy as np + +from _example_paths import TEST_DATA, prepare_output_dir + +DATA = TEST_DATA / "generated" / "2d_c2p_rot45_ms_p1.gkyl" +OUTPUT_DIR = prepare_output_dir() + +# 1. Load -- raw DG coefficients, native to Gkeyll. This 2D file carries a +# 2-component vector field (a rotated coordinate map). +d = pg.load(DATA) +print("loaded:", repr(d)) + +# 2. Interpolate -- the modal -> NumPy bridge. +field = d.interpolate() +print("interpolated:", repr(field)) + +# 3a. Select by component -- pick out the x-component as its own 2D scalar +# field, and plot it. +comp0 = field.select(comp=0) +fig = comp0.plot(title="x-component", no_show=True) +png_path = OUTPUT_DIR / "01_quickstart_comp0.png" +fig.savefig(png_path) +print("saved plot:", png_path) + +# 3b. Select by coordinate -- fix y=0.5 to take a lineout across x, keeping +# both components. The y axis survives as a size-1 slice (``select`` +# narrows a coordinate, it doesn't drop the axis -- ``plot`` squeezes it +# away when rendering). +lineout = field.select(z1=0.5) +fig = lineout.plot(title="lineout at y=0.5", no_show=True) +lineout_path = OUTPUT_DIR / "01_quickstart_lineout.png" +fig.savefig(lineout_path) +print("saved plot:", lineout_path) + +# 4. Save / reload round trip -- writing an interpolated field to ``.gkyl`` +# and reading it back recovers the same values exactly. +gkyl_path = comp0.save(str(OUTPUT_DIR / "01_quickstart_out.gkyl")) +reloaded = pg.load(gkyl_path) +np.testing.assert_allclose(reloaded.values, comp0.values, rtol=0, atol=0) +print("round-tripped through", gkyl_path) + +print("01_quickstart: OK") diff --git a/examples/scripts/02_arithmetic_and_numpy.py b/examples/scripts/02_arithmetic_and_numpy.py new file mode 100644 index 00000000..374b9465 --- /dev/null +++ b/examples/scripts/02_arithmetic_and_numpy.py @@ -0,0 +1,83 @@ +"""Weak (DG) arithmetic and NumPy interop. + +This example contrasts two representations: + +* **modal** (straight off disk): ``+``/``-``/``*``/``/`` on two ``GData`` + run *inside Gkeyll* (coefficient lin-combs / weak DG kernels), and + ``.integrate()`` runs a grid integral there too. Plain NumPy math + (``np.sqrt``, ``np.asarray``, ...) refuses -- coefficients aren't values. +* **numpy** (after ``.interpolate()``): the field is plain point values, so + every NumPy ufunc and ``+``/``-``/``*``/``/`` "just work", and the result + keeps carrying its grid/``ctx`` like a ``GData`` should. + +Run directly: + MPLBACKEND=Agg PYTHONPATH=src python examples/scripts/02_arithmetic_and_numpy.py +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") + +import numpy as np + +import postgkyl as pg + +from _example_paths import TEST_DATA, prepare_output_dir + +# A coordinate-map field: strictly positive everywhere, so weak division +# never divides near zero (see the note on ``back`` below). +DATA = TEST_DATA / "generated" / "2d_c2p_stretch_ms_p1.gkyl" + +# ----------------------------------------------------------- modal algebra +a = pg.load(DATA) +b = pg.load(DATA) + +prod = a * b # gkyl_dg_mul_op -- weak multiply +back = prod / b # gkyl_dg_div_op -- weak divide +summ = a + b # gkyl_array_accumulate -- coefficient sum + +# ``a + a == 2*a`` is an exact linear identity -- coefficient lin-combs never +# alias. Weak multiply/divide, by contrast, is a *nonlinear* operation: the +# product of two degree-p polynomials is degree 2p, and projecting it back +# onto the degree-p basis is lossy in general. It only round-trips exactly +# for special inputs. A nonzero divisor does not by itself guarantee an +# exact round trip. Compare via NumPy, which means interpolating first. +print("(a*b)/b == a: ", + np.allclose(back.interpolate().values, + a.interpolate().values)) +print("a+a == 2*a: ", + np.allclose(summ.interpolate().values, (2.0 * a).interpolate().values)) +np.testing.assert_allclose(summ.interpolate().values, + (2.0 * a).interpolate().values) + +total = a.integrate() # gkyl_array_integrate -- a terminal verb +print("integrate(a) = ", total) # one value per component (this field has 2) + +# A general NumPy ufunc has no meaning on raw DG coefficients -- the +# capability boundary is enforced, not just "usually correct". +try: + np.sqrt(a) +except ValueError as exc: + print("np.sqrt(modal) refuses:", exc) + +# ----------------------------------------------------- field-domain (NumPy) +fa = a.interpolate() +fb = b.interpolate() + +field_sum = fa + fb +mag = np.sqrt(fa**2 + fb**2) # ufunc -> still a GData, carrying fa's grid +print("sqrt(a^2+b^2) == sqrt(2)*a:", + np.allclose(np.asarray(mag), + np.sqrt(2) * np.asarray(fa))) + +as_array = np.asarray(fa) # plain ndarray -- the escape hatch out of GData +print("np.asarray(interpolated) ->", as_array.shape, as_array.dtype) + +mag.plot(title="Componentwise magnitude", + no_show=True, + saveas=prepare_output_dir() / "02_arithmetic.png", + dpi=100) + +print("02_arithmetic_and_numpy: OK") diff --git a/examples/scripts/03_diagnostics_five_moment.py b/examples/scripts/03_diagnostics_five_moment.py new file mode 100644 index 00000000..2a119e92 --- /dev/null +++ b/examples/scripts/03_diagnostics_five_moment.py @@ -0,0 +1,67 @@ +"""The diagnostics layer: equation-specific physics on top of the same +``GData`` you get from ``pg.load``. + +``diagnostics`` sits *above* the fluent API (``operations``/``api``) and is +equation-blind-free by design: ``postgkyl.diagnostics.mom.five_moment`` knows the +Euler fluid moment layout (``[rho, rho*vx, rho*vy, rho*vz, E]``) and turns raw +conserved moments into primitive variables (density, velocity, pressure, +Mach number, ...). Diagnostics are **free functions**, not ``GData`` methods +-- ``fm.density(d)``, never ``d.density()`` -- because a diagnostic knows +about one specific equation system and a ``GData`` doesn't. + +This script loads a generated stationary shock-tube initial condition and +checks density, pressure, and Mach number against the prescribed state. +Run ``python tests/generate_test_data.py`` first. + +Run directly: + MPLBACKEND=Agg PYTHONPATH=src python examples/scripts/03_diagnostics_five_moment.py +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") + +import numpy as np + +import postgkyl as pg +from postgkyl.diagnostics.mom import five_moment as fm + +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() + +GAS_GAMMA = 5.0 / 3.0 + +# Load the generated p0 modal state and evaluate its physical values. +d = pg.load(TEST_DATA / "generated/shock_tube_1d_p0.gkyl").interpolate() +x = 0.5 * (d.grid[0][1:] + d.grid[0][:-1]) +rho = np.where(x < 0.5, 1.0, 0.125) +p = np.where(x < 0.5, 1.0, 0.1) + +# Diagnostics are free functions of a GData(State), returning a new one -- +# the same ``inplace``/``tag``/``label`` contract as every ``operations`` verb. +density = fm.density(d) +pressure = fm.pressure(d, gas_gamma=GAS_GAMMA) +mach = fm.mach(d, gas_gamma=GAS_GAMMA) + +np.testing.assert_allclose(density.values.ravel(), rho) +np.testing.assert_allclose(pressure.values.ravel(), p) +np.testing.assert_allclose(mach.values, 0.0) + +# Raw modal DG coefficients have no "density"/"pressure" until interpolated +# -- the diagnostics layer refuses the same way ``operations`` verbs do. +modal = pg.load(TEST_DATA / "generated" / "1d_ms_p1.gkyl") +try: + fm.density(modal) +except ValueError as exc: + print("fm.density(modal) refuses:", exc) + +fig = density.plot(title="density (Sod shock tube)", no_show=True) +fig.savefig(OUTPUT_DIR / "03_diagnostics_density.png") + +fig = pressure.plot(title="pressure (Sod shock tube)", no_show=True) +fig.savefig(OUTPUT_DIR / "03_diagnostics_pressure.png") + +print("03_diagnostics_five_moment: OK") diff --git a/examples/scripts/04_gyrokinetics.py b/examples/scripts/04_gyrokinetics.py new file mode 100644 index 00000000..3fcb68f3 --- /dev/null +++ b/examples/scripts/04_gyrokinetics.py @@ -0,0 +1,94 @@ +"""Gyrokinetic diagnostics: named quantities and distribution functions. + +``postgkyl.diagnostics.gk`` is another equation-specific module in +the diagnostics layer (like ``five_moment`` in +``examples/scripts/03_diagnostics_five_moment.py``), but for gyrokinetic +simulations it owns its own *loading* too -- a simulation's files follow a +naming convention (``-__.gkyl``), so instead +of calling ``pg.load`` on individual filenames, you ask for a quantity by +name and the loader resolves which files it needs: + +* ``pg.gk.available_quantities()`` lists the registered quantity names. +* ``pg.gk.load_quantity(quantity, species, name, frame, path=...)`` resolves + the source file(s) for that quantity, computes it, and returns one + ``GData`` per requested species -- already interpolated, ready to + ``.select()``/``.plot()`` like anything else. +* ``pg.gk.load_distf(name=..., species=..., frame=..., ...)`` reconstructs a + full distribution function from the saved ``Jf``-times-Jacobian(s) files + (what the CLI's ``gk_load_distf`` command wraps). + +This uses the ``rt_gk_tcv_iwl*`` fixtures staged in ``tests/test_data`` -- +two related simulations: ``rt_gk_tcv_iwl_adapt_source_1x2v_p1`` wrote ion +Hamiltonian moments (``M0``/``M1`` are derivable from those), and +``rt_gk_tcv_iwl_1x2v_p1`` wrote the electron distribution function plus its +geometry (``geo_int_jacobtot_inv``). + +Run directly: + MPLBACKEND=Agg PYTHONPATH=src python examples/scripts/04_gyrokinetics.py +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg + +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() + +HMOM_NAME = "rt_gk_tcv_iwl_adapt_source_1x2v_p1" # wrote ion Hamiltonian moments +GK_NAME = "rt_gk_tcv_iwl_1x2v_p1" # wrote elc distf + geometry + +print("registered quantities:", pg.gk.available_quantities()) + +# 1. A moment quantity ("M0", the density) resolved straight from the +# HamiltonianMoments file the registry knows how to read; already +# interpolated, so it's a regular one-component field. +m0, = pg.gk.load_quantity("M0", "ion", HMOM_NAME, "250", path=str(TEST_DATA)) +print("M0:", repr(m0), " label:", m0.get_label()) + +fig = m0.plot(title="Ion density from Hamiltonian moments", + xlabel="Field-aligned computational coordinate", + ylabel=m0.get_label(), + no_show=True) +fig.savefig(OUTPUT_DIR / "04_gyrokinetics_M0.png") + +# M1 is a density-weighted parallel velocity moment, not velocity itself. +# Its Hamiltonian conversion requires the simulation's species mass in the +# same unit system. This fixture has no input deck establishing that mass, +# so this example limits itself to M0 rather than assigning an arbitrary mass. + +# 3. A species-independent geometric factor, from the other simulation -- +# ``species=None`` since geometry isn't per-species. +jacobtot_inv, = pg.gk.load_quantity("geo_int_jacobtot_inv", + None, + GK_NAME, + path=str(TEST_DATA)) +print("(J B)^-1:", repr(jacobtot_inv), " label:", jacobtot_inv.get_label()) + +fig = jacobtot_inv.plot(title=jacobtot_inv.get_label(), no_show=True) +fig.savefig(OUTPUT_DIR / "04_gyrokinetics_jacobtot_inv.png") + +# 4. The full electron distribution function: 3D (x, vpar, mu), built from +# the saved Jf-times-Jacobian(s) file plus the geometry factor above. +distf = pg.gk.load_distf(name=str(TEST_DATA / GK_NAME), + species="elc", + frame=250, + jacobtot_inv_file=TEST_DATA / + f"{GK_NAME}-geo_int_jacobtot_inv.gkyl") +print("distf:", repr(distf)) + +# It's a regular GData from here on -- e.g. select a fixed-mu slice down to +# the (x, vpar) plane and plot it, same as any other 2D field. +slice_2d = distf.select(z2=0.0) +fig = slice_2d.plot(title="Electron distribution: lowest mu slice", + xlabel="Field-aligned computational coordinate", + ylabel="Computational parallel velocity", + clabel="Distribution [simulation units]", + no_show=True) +fig.savefig(OUTPUT_DIR / "04_gyrokinetics_distf_slice.png") + +print("04_gyrokinetics: OK") diff --git a/examples/scripts/05_gk_rz.py b/examples/scripts/05_gk_rz.py new file mode 100644 index 00000000..40927921 --- /dev/null +++ b/examples/scripts/05_gk_rz.py @@ -0,0 +1,50 @@ +"""Gyrokinetic R-Z mapping through fluent, functional, and reusable APIs. + +Run directly: + MPLBACKEND=Agg PYTHONPATH=src python examples/scripts/05_gk_rz.py +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import numpy as np + +import postgkyl as pg +from postgkyl.operations import gyrokinetics as gk_ops + +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() + +data = pg.load(TEST_DATA / "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") + +# The common path: geometry is inferred from FIELD's simulation prefix. +mapped = data.gk_rz(z_axis=0.0, phi_tor=0.0, nz_interp=2) +fig = mapped.plot(title="Electron density in the poloidal plane", + xlabel="R [m]", + ylabel="Z [m]", + clabel=r"$n_e$ [m$^{-3}$]", + fixaspect=True, + no_show=True) +fig.savefig(OUTPUT_DIR / "05_gk_rz.png") + +# The functional spelling is the identical operation. +functional = pg.gk_rz(data, z_axis=0.0, phi_tor=0.0, nz_interp=2) +np.testing.assert_allclose(functional.values, mapped.values) + +# Reuse geometry and projection when several fields/frames or toroidal angles +# share one computational grid. +geometry = gk_ops.resolve_geometry(data.file_name) +projection = gk_ops.resolve_rz_projection(data, + geometry, + z_axis=0.0, + nz_interp=2) +at_zero = gk_ops.map_to_rz(data, projection, phi_tor=0.0) +at_quarter_turn = gk_ops.map_to_rz(data, projection, phi_tor=np.pi / 2) +np.testing.assert_allclose(at_zero.values, mapped.values) +assert not np.allclose( + at_zero.values, at_quarter_turn.values, rtol=1e-12, atol=0.0) + +print("05_gk_rz: OK") diff --git a/examples/scripts/06_growth.py b/examples/scripts/06_growth.py new file mode 100644 index 00000000..0b8e4238 --- /dev/null +++ b/examples/scripts/06_growth.py @@ -0,0 +1,27 @@ +"""Recover a known energy growth rate from a generated diagnostic history.""" + +import matplotlib + +matplotlib.use("Agg") +import numpy as np +import postgkyl as pg + +from _example_paths import TEST_DATA, prepare_output_dir + +data = pg.load(TEST_DATA / "generated/exponential_energy.gkyl") +log_energy = pg.evaluate("f log", data) +fit = log_energy.fit("linear") +time = data.grid[0] +energy_rate = np.polyfit(time, fit.values.ravel(), 1)[0] +np.testing.assert_allclose(energy_rate, 0.4, rtol=1e-10) +np.testing.assert_allclose(fit.values, log_energy.values, atol=1e-12) + +# Plot the reconstructed fitted energy with the public plotting function. +# The CLI uses evaluate "f log", fit linear, and evaluate "f1 exp" (fit appends its result). +pg.evaluate("f exp", fit).plot(title="Fitted energy growth", + xlabel="Time [normalized]", + ylabel="Energy [normalized]", + logy=True, + no_show=True, + saveas=prepare_output_dir() / "06_growth.png", + dpi=100) diff --git a/examples/scripts/07_collect_animate.py b/examples/scripts/07_collect_animate.py new file mode 100644 index 00000000..10ee7ce8 --- /dev/null +++ b/examples/scripts/07_collect_animate.py @@ -0,0 +1,43 @@ +"""Load numbered frames, assemble an x–t diagram, and animate 1D and 2D data.""" + +import numpy as np +import postgkyl as pg +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() +paths = sorted((TEST_DATA / "generated").glob("travelling_wave_*.gkyl")) +frames = [pg.load(path).interpolate() for path in paths] +assert len(frames) == 16 + +# collect adds a leading time coordinate, taken from each file's metadata. +history = pg.collect(frames) +np.testing.assert_allclose(history.grid[0], [d.ctx["time"] for d in frames]) +np.testing.assert_allclose(history.values, np.stack([d.values for d in frames])) +history.plot(title="Travelling wave: space–time diagram", + xlabel="Time", + ylabel="x", + clabel="Density [normalized]", + no_show=True, + saveas=OUTPUT_DIR / "07_collect.png", + dpi=100) + +# animate takes the individual frames, not the dataset returned by collect. +# saveframes also keeps every PNG for analysis or a presentation. +pg.animate(frames, + saveframes=str(OUTPUT_DIR / "07_wave_frame"), + saveas=str(OUTPUT_DIR / "07_wave.gif"), + fps=8, + dpi=90, + figsize=(6, 4), + no_show=True) + +surface_paths = sorted((TEST_DATA / "generated").glob("wave_surface_*.gkyl")) +surface_frames = [pg.load(path).interpolate() for path in surface_paths] +assert len(surface_frames) == 12 +pg.animate(surface_frames, + saveframes=str(OUTPUT_DIR / "07_surface_frame"), + saveas=str(OUTPUT_DIR / "07_surface.gif"), + fps=6, + dpi=90, + figsize=(6, 4), + no_show=True) diff --git a/examples/scripts/08_plotly.py b/examples/scripts/08_plotly.py new file mode 100644 index 00000000..16e02587 --- /dev/null +++ b/examples/scripts/08_plotly.py @@ -0,0 +1,27 @@ +"""Interactive height surfaces from 2D data and isosurfaces from a 3D field.""" + +import postgkyl as pg +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() +surface = pg.load(TEST_DATA / "generated/wave_surface_000.gkyl").interpolate() +surface.plotly(background="light", + title="Travelling-wave height surface", + xlabel="x", + ylabel="y", + zlabel="Density [normalized]", + clabel="Density", + saveas=str(OUTPUT_DIR / "08_surface.html")) + +volume = pg.load(TEST_DATA / "generated/gaussian_volume.gkyl").interpolate() +volume.plotly(background="light", + title="Gaussian density isosurfaces", + xlabel="x", + ylabel="y", + zlabel="z", + clabel="Density", + surface_count=8, + opacity=0.15, + cmin=0.1, + cmax=0.9, + saveas=str(OUTPUT_DIR / "08_volume.html")) diff --git a/examples/scripts/09_pyvista.py b/examples/scripts/09_pyvista.py new file mode 100644 index 00000000..5235f138 --- /dev/null +++ b/examples/scripts/09_pyvista.py @@ -0,0 +1,34 @@ +"""Render Gaussian isosurfaces and a translucent volume with PyVista.""" + +import postgkyl as pg +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() +data = pg.load(TEST_DATA / "generated/gaussian_volume.gkyl").interpolate() +data.pyvista(contour_levels=5, + no_show=True, + no_spin=True, + hide_axes=True, + title="Gaussian density isosurfaces", + xlabel="x", + ylabel="y", + zlabel="z", + clabel="Density", + theme="document", + camera_elevation=10.0, + camera_azimuth=20.0, + saveas=str(OUTPUT_DIR / "09_isosurfaces.png")) +data.pyvista(volume=True, + opacity="linear", + no_show=True, + no_spin=True, + hide_axes=True, + title="Gaussian density volume", + xlabel="x", + ylabel="y", + zlabel="z", + clabel="Density", + theme="document", + camera_elevation=10.0, + camera_azimuth=20.0, + saveas=str(OUTPUT_DIR / "09_volume.png")) diff --git a/examples/scripts/10_manual_arrays.py b/examples/scripts/10_manual_arrays.py new file mode 100644 index 00000000..1c41005c --- /dev/null +++ b/examples/scripts/10_manual_arrays.py @@ -0,0 +1,44 @@ +"""Copy grids and values, manipulate the arrays, and wrap a field for plotting.""" + +import numpy as np +import postgkyl as pg +from _example_paths import TEST_DATA, prepare_output_dir + +OUTPUT_DIR = prepare_output_dir() +modal = pg.load(TEST_DATA / "generated/wave_surface_000.gkyl") +# On modal data, these numbers are coefficients, not physical field samples. +coefficients = modal.values.copy() +assert coefficients.shape[-1] == 1 # this generated example uses p0 + +field = modal.interpolate() +grid = [axis.copy() for axis in field.grid] +values = field.values.copy() +np.testing.assert_array_equal(np.asarray(field), values) + +# This interpolated mesh has edges. Pair its cell centers with field values; +# for data already located at nodes, use the grid coordinates directly. +coordinates = [ + 0.5 * (axis[1:] + axis[:-1]) if len(axis) == values.shape[dim] + 1 else axis + for dim, axis in enumerate(grid) +] +x, y = np.meshgrid(*coordinates, indexing="ij") +assert x.shape == y.shape == values.shape[:-1] + +# NumPy edits to these copies do not change the original dataset. +manual_values = (values - 1.0) / 0.6 +np.testing.assert_allclose(manual_values[..., 0], + np.cos(x) * np.cos(y), + atol=1e-14) +np.testing.assert_array_equal(field.values, values) + +# push accepts coordinates and values. Preserve metadata explicitly when +# constructing a new point-value dataset after a manual calculation. +manual = field.clone() +manual.push(grid, manual_values) +manual.plot(title="Normalized density perturbation", + xlabel="x", + ylabel="y", + clabel="(n - 1) / 0.6", + no_show=True, + saveas=OUTPUT_DIR / "10_manual.png", + dpi=100) diff --git a/examples/scripts/_example_paths.py b/examples/scripts/_example_paths.py new file mode 100644 index 00000000..9a8e3ef7 --- /dev/null +++ b/examples/scripts/_example_paths.py @@ -0,0 +1,19 @@ +"""Shared filesystem locations for the executable examples.""" + +from __future__ import annotations + +import os +from pathlib import Path + +_SCRIPT_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SCRIPT_DIR.parents[1] + +TEST_DATA = _REPO_ROOT / "tests" / "test_data" + + +def prepare_output_dir() -> Path: + """Return the configured output directory, creating it if necessary.""" + output_dir = Path( + os.environ.get("PGKYL_EXAMPLE_OUTPUT", _SCRIPT_DIR / "output")) + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir diff --git a/examples/scripts/mirror_comparison.py b/examples/scripts/mirror_comparison.py new file mode 100644 index 00000000..20923897 --- /dev/null +++ b/examples/scripts/mirror_comparison.py @@ -0,0 +1,37 @@ +import postgkyl as pg +from _example_paths import TEST_DATA, prepare_output_dir + +DATA = TEST_DATA / "generated" +OUTPUT_DIR = prepare_output_dir() +figure_path = OUTPUT_DIR / "mirror_comparison.pdf" + +high_alpha = pg.load(DATA / + "mirror_comparison_2em4_1d_ms_p1.gkyl").interpolate() +reference = pg.load(DATA / "mirror_comparison_2em5_1d_ms_p1.gkyl").interpolate() + +fig = pg.plot( + high_alpha, + reference, + figure=0, + legend_labels=["2e-4", "2e-5"], + no_legend=False, + legend_subplot=0, + legend_loc="best", + title=r"Comparing $\alpha$ = 2e-4 to $\alpha$ = 2e-5", + xlabel="z [m]", + figsize="12,10", + split_linear_log=True, + split_point=0.0, + split_log_side="right", + split_width_ratios=(1.0, 1.0), + split_gap=0.0, + split_legend_side="log", + split_log_nonpositive="mask", + subplot_ylabels=(r"Density [$\mathrm{m}^{-3}$]," + r"$U_\parallel$ [m/s]," + r"$T_\parallel$ [keV]," + r"$T_\perp$ [keV]"), + saveas=figure_path, + no_show=True, +) +fig.savefig(OUTPUT_DIR / "mirror_comparison.png", dpi=150) diff --git a/gkeyll-data/.gitignore b/gkeyll-data/.gitignore new file mode 100644 index 00000000..4b5a10a7 --- /dev/null +++ b/gkeyll-data/.gitignore @@ -0,0 +1 @@ +*.gkyl diff --git a/notebooks/README.md b/notebooks/README.md index 07adab1a..2187a22a 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -1,7 +1,7 @@ -postgkyl notebooks +postgkyl notebooks ================== -Notebooks using Marimo or Jupyter are useful collections of operations to pre- or post-process Gkeyll simulations, or better understand their algorithms. This folder contains some such notebooks. +Notebooks using Marimo or Jupyter are useful collections of operations to pre- or post-process Gkeyll simulations, or better understand their algorithms. This folder contains some such notebooks. File convention: We use the `*.mo.py` extension for Marimo notebooks, which are Python scripts that can be executed in a notebook-like environment. - twist_shift.mo.py: Analyze the impact of twist-shift boundary conditions on mode shearing and aliasing, and test possible filtering strategies to mitigate these effects. diff --git a/notebooks/dg_twist_shift.mo.py b/notebooks/dg_twist_shift.mo.py index 21cc838d..08dacd02 100644 --- a/notebooks/dg_twist_shift.mo.py +++ b/notebooks/dg_twist_shift.mo.py @@ -13,6 +13,7 @@ # - gkeyll libraries (core and gyrokinetic). # - Marimo. + @app.cell def _imports(): import glob @@ -29,8 +30,17 @@ def _imports(): import postgkyl as pg import marimo as mo return ( - glob, re, subprocess, Path, - matplotlib, plt, patches, Line2D, np, pg, mo, + glob, + re, + subprocess, + Path, + matplotlib, + plt, + patches, + Line2D, + np, + pg, + mo, ) @@ -45,9 +55,10 @@ def getRawGrid(dataFile, **opKey): if opKey['location'] == 'center': xOut = [[] for _ in range(dimOut)] for i in range(dimOut): - nNodes = np.shape(xNodal[i])[0] + nNodes = np.shape(xNodal[i])[0] xOut[i] = np.zeros(nNodes - 1) - xOut[i] = np.multiply(0.5, xNodal[i][0:nNodes-1] + xNodal[i][1:nNodes]) + xOut[i] = np.multiply(0.5, + xNodal[i][0:nNodes - 1] + xNodal[i][1:nNodes]) else: xOut = xNodal else: @@ -74,85 +85,91 @@ def getRawData(dataFile): def _controls(mo): # Compile / run controls sim_dir = mo.ui.text( - value="/Users/mfrancis/Documents/gkeyll/code/gkeyll_v0/cbc_tst", - label="Simulation directory", - full_width=True, + value="/Users/mfrancis/Documents/gkeyll/code/gkeyll_v0/cbc_tst", + label="Simulation directory", + full_width=True, ) input_file_ui = mo.ui.text( - value="rt_gk_cbc_passive_3x2v_p1", - label="Input file name", - full_width=True, + value="rt_gk_cbc_passive_3x2v_p1", + label="Input file name", + full_width=True, ) sim_name_ui = mo.ui.text( - value="rt_gk_cbc_passive_3x2v_p1", - label="Simulation name", - full_width=True, + value="rt_gk_cbc_passive_3x2v_p1", + label="Simulation name", + full_width=True, ) gkylsoft_ui = mo.ui.text( - value=str(__import__('pathlib').Path.home() / "gkylsoft"), - label="gkylsoft path", - full_width=True, + value=str(__import__('pathlib').Path.home() / "gkylsoft"), + label="gkylsoft path", + full_width=True, ) - ly_fac = mo.ui.number(value=1, step=0.1, label="Ly factor") - nx = mo.ui.number(value=32, step=1, label="Nx") - ny = mo.ui.number(value=16, step=1, label="Ny") + ly_fac = mo.ui.number(value=1, step=0.1, label="Ly factor") + nx = mo.ui.number(value=32, step=1, label="Nx") + ny = mo.ui.number(value=16, step=1, label="Ny") run_btn = mo.ui.run_button(label="Compile & Run") # Field-plot controls - species = mo.ui.text(value="elc", label="Species") - quantity = mo.ui.text(value="M0", label="Quantity") - comp_ui = mo.ui.number(value=0, step=1, label="Component") + species = mo.ui.text(value="elc", label="Species") + quantity = mo.ui.text(value="M0", label="Quantity") + comp_ui = mo.ui.number(value=0, step=1, label="Component") boundary = mo.ui.dropdown( - options={"Lower z": "lower", "Upper z": "upper"}, - value="Lower z", - label="Boundary", + options={ + "Lower z": "lower", + "Upper z": "upper" + }, + value="Lower z", + label="Boundary", ) return sim_dir, input_file_ui, sim_name_ui, gkylsoft_ui, ly_fac, nx, ny, run_btn, species, quantity, comp_ui, boundary @app.cell -def _layout(mo, sim_dir, input_file_ui, sim_name_ui, gkylsoft_ui, ly_fac, nx, ny, run_btn, - species, quantity, comp_ui, boundary): +def _layout(mo, sim_dir, input_file_ui, sim_name_ui, gkylsoft_ui, ly_fac, nx, + ny, run_btn, species, quantity, comp_ui, boundary): mo.vstack([ - sim_dir, - input_file_ui, - sim_name_ui, - gkylsoft_ui, - mo.hstack([ly_fac, nx, ny, run_btn], justify="start"), - mo.callout( - mo.md( - "**Assumptions:** the input C file specifies the number of cells as `Nx` and `Ny`, " - "and the domain size along y as `Ly` (in units of `rho_s`)." + sim_dir, + input_file_ui, + sim_name_ui, + gkylsoft_ui, + mo.hstack([ly_fac, nx, ny, run_btn], justify="start"), + mo.callout( + mo. + md("**Assumptions:** the input C file specifies the number of cells as `Nx` and `Ny`, " + "and the domain size along y as `Ly` (in units of `rho_s`)."), + kind="info", ), - kind="info", - ), - mo.md("---"), - mo.md("**Field plot**"), - mo.hstack([species, quantity, comp_ui, boundary], justify="start"), + mo.md("---"), + mo.md("**Field plot**"), + mo.hstack([species, quantity, comp_ui, boundary], justify="start"), ]) @app.cell -def _compile_run(mo, re, subprocess, Path, - sim_dir, input_file_ui, sim_name_ui, ly_fac, nx, ny, run_btn): +def _compile_run(mo, re, subprocess, Path, sim_dir, input_file_ui, sim_name_ui, + ly_fac, nx, ny, run_btn): # Only execute when the run button is pressed. mo.stop(not run_btn.value, mo.md("Click **Compile & Run** to start.")) - _sdir = sim_dir.value.rstrip("/") - _input_file = input_file_ui.value.strip() - _sim_name = sim_name_ui.value.strip() + _sdir = sim_dir.value.rstrip("/") + _input_file = input_file_ui.value.strip() + _sim_name = sim_name_ui.value.strip() # Find the C input file _c_path = Path(f"{_sdir}/{_input_file}.c") if not _c_path.exists(): - mo.stop(True, mo.callout(mo.md(f"**Error:** `{_c_path}` not found"), kind="danger")) + mo.stop( + True, + mo.callout(mo.md(f"**Error:** `{_c_path}` not found"), kind="danger")) # Warn if symlink if _c_path.is_symlink(): mo.callout( - mo.md(f"⚠️ `{_c_path.name}` is a symlink — edits will also modify the original file."), - kind="warn", + mo. + md(f" `{_c_path.name}` is a symlink -- edits will also modify the original file." + ), + kind="warn", ) # Backup original on first run @@ -165,35 +182,39 @@ def _compile_run(mo, re, subprocess, Path, _content = _c_path.read_text() _content = re.sub( - r'(double Ly\s*=\s*)([\d.]+)(\*rho_s\s*;)', - lambda m: f"{m.group(1)}{float(m.group(2)) * ly_fac.value}{m.group(3)}", - _content) - _content = re.sub( - r'(int Nx\s*=\s*)\d+(\s*;)', - rf'\g<1>{int(nx.value)}\g<2>', _content) - _content = re.sub( - r'(int Ny\s*=\s*)\d+(\s*;)', - rf'\g<1>{int(ny.value)}\g<2>', _content) + r'(double Ly\s*=\s*)([\d.]+)(\*rho_s\s*;)', + lambda m: f"{m.group(1)}{float(m.group(2)) * ly_fac.value}{m.group(3)}", + _content) + _content = re.sub(r'(int Nx\s*=\s*)\d+(\s*;)', rf'\g<1>{int(nx.value)}\g<2>', + _content) + _content = re.sub(r'(int Ny\s*=\s*)\d+(\s*;)', rf'\g<1>{int(ny.value)}\g<2>', + _content) _c_path.write_text(_content) # Compile - _make = subprocess.run( - ['make', _input_file], cwd=_sdir, - capture_output=True, text=True) + _make = subprocess.run(['make', _input_file], + cwd=_sdir, + capture_output=True, + text=True) if _make.returncode != 0: _msg = _make.stdout + "\n" + _make.stderr - mo.stop(True, mo.callout(mo.md(f"**make failed:**\n```\n{_msg}\n```"), kind="danger")) + mo.stop( + True, + mo.callout(mo.md(f"**make failed:**\n```\n{_msg}\n```"), kind="danger")) # Run initialization - _run = subprocess.run( - [f'./{_input_file}', '-s0'], cwd=_sdir, - capture_output=True, text=True) + _run = subprocess.run([f'./{_input_file}', '-s0'], + cwd=_sdir, + capture_output=True, + text=True) if _run.returncode != 0: _msg = _run.stdout + "\n" + _run.stderr - mo.stop(True, mo.callout(mo.md(f"**Run failed:**\n```\n{_msg}\n```"), kind="danger")) + mo.stop( + True, + mo.callout(mo.md(f"**Run failed:**\n```\n{_msg}\n```"), kind="danger")) result = (True, _sim_name, _sdir) return result, @@ -208,7 +229,7 @@ def _donor_plot(mo, result, getRawGrid, getRawData, np, plt, patches, Line2D): _ok, _sim_name, _sdir = result # Load grid and shift from .gkyl output files - _geo_file = f"{_sdir}/{_sim_name}-geo_corn_bmag.gkyl" + _geo_file = f"{_sdir}/{_sim_name}-geo_corn_bmag.gkyl" _shift_file = f"{_sdir}/{_sim_name}-bc_zlower_twistshift.gkyl" _x_nodal, _, _geo_nx, _, _, _ = getRawGrid(_geo_file) @@ -218,18 +239,18 @@ def _donor_plot(mo, result, getRawGrid, getRawData, np, plt, patches, Line2D): _x_nodes = _x_nodal[0] _y_nodes = _x_nodal[1] - _x_min, _x_max = _x_nodes[0], _x_nodes[-1] - _y_min, _y_max = _y_nodes[0], _y_nodes[-1] + _x_min, _x_max = _x_nodes[0], _x_nodes[-1] + _y_min, _y_max = _y_nodes[0], _y_nodes[-1] _Lx, _Ly = _x_max - _x_min, _y_max - _y_min _dx, _dy = _Lx / _Nx, _Ly / _Ny _x_centers = 0.5 * (_x_nodes[:-1] + _x_nodes[1:]) - _dx_cells = np.diff(_x_nodes) + _dx_cells = np.diff(_x_nodes) - _shift_data = getRawData(_shift_file) # shape (Nx, 2) + _shift_data = getRawData(_shift_file) # shape (Nx, 2) # Shift evaluation (p=1 DG) - _sqrt2 = np.sqrt(2.0) + _sqrt2 = np.sqrt(2.0) _sqrt32 = np.sqrt(1.5) def _S_eval(x_arr): @@ -240,7 +261,7 @@ def _S_eval(x_arr): _S_lo_cells = _shift_data[:, 0] / _sqrt2 - _sqrt32 * _shift_data[:, 1] _S_up_cells = _shift_data[:, 0] / _sqrt2 + _sqrt32 * _shift_data[:, 1] - _dS_cells = _S_up_cells - _S_lo_cells + _dS_cells = _S_up_cells - _S_lo_cells def _wrap(v): return _y_min + np.mod(v - _y_min, _Ly) @@ -257,15 +278,20 @@ def _wrap(v): # Grid for _ix in range(_Nx): for _iy in range(_Ny): - _ax.add_patch(patches.Rectangle( - (_x_nodes[_ix], _y_nodes[_iy]), _dx_cells[_ix], _dy, - lw=0.4, edgecolor='#aaaaaa', facecolor='white', zorder=1)) + _ax.add_patch( + patches.Rectangle((_x_nodes[_ix], _y_nodes[_iy]), + _dx_cells[_ix], + _dy, + lw=0.4, + edgecolor='#aaaaaa', + facecolor='white', + zorder=1)) # Donor cells for _ix in range(_Nx): _x_lo = _x_nodes[_ix] _x_up = _x_nodes[_ix + 1] - _xs = np.linspace(_x_lo, _x_up, 500) + _xs = np.linspace(_x_lo, _x_up, 500) _donor_iys = set() for _frac in np.linspace(0.01, 0.99, 9): _y_probe = _y_tar_lo + _frac * _dy @@ -273,62 +299,102 @@ def _wrap(v): _iys = np.clip(np.floor((_yd - _y_min) / _dy).astype(int), 0, _Ny - 1) _donor_iys.update(np.unique(_iys)) for _iy in _donor_iys: - _ax.add_patch(patches.Rectangle( - (_x_lo, _y_nodes[_iy]), _dx_cells[_ix], _dy, - alpha=0.38, facecolor='royalblue', edgecolor='none', zorder=2)) + _ax.add_patch( + patches.Rectangle((_x_lo, _y_nodes[_iy]), + _dx_cells[_ix], + _dy, + alpha=0.38, + facecolor='royalblue', + edgecolor='none', + zorder=2)) # Wrapped boundary lines _x_dense = np.linspace(_x_min, _x_max, 8000) - _S_dense = _S_eval(_x_dense) + _S_dense = _S_eval(_x_dense) def _plot_wrapped_line(y_raw, ls, lbl): - y = _wrap(y_raw) + y = _wrap(y_raw) jump = np.abs(np.diff(y)) > _Ly / 2 mask = np.concatenate([[False], jump]) | np.concatenate([jump, [False]]) - _ax.plot(_x_dense, np.where(mask, np.nan, y), - color='navy', lw=1.8, ls=ls, zorder=5, label=lbl) - - _plot_wrapped_line(_y_tar_lo - _S_dense, '-', r'$y_{\rm tar,lo} - S(x)$') + _ax.plot(_x_dense, + np.where(mask, np.nan, y), + color='navy', + lw=1.8, + ls=ls, + zorder=5, + label=lbl) + + _plot_wrapped_line(_y_tar_lo - _S_dense, '-', r'$y_{\rm tar,lo} - S(x)$') _plot_wrapped_line(_y_tar_up - _S_dense, '--', r'$y_{\rm tar,up} - S(x)$') # Wrap points for _ix in range(_Nx): - _x_lo = _x_nodes[_ix]; _x_up = _x_nodes[_ix + 1] - _S_lo = _S_lo_cells[_ix]; _S_up = _S_up_cells[_ix] + _x_lo = _x_nodes[_ix] + _x_up = _x_nodes[_ix + 1] + _S_lo = _S_lo_cells[_ix] + _S_up = _S_up_cells[_ix] if _S_lo == _S_up: continue - _n_lo = int(np.ceil( _S_lo / _Ly)) + _n_lo = int(np.ceil(_S_lo / _Ly)) _n_up = int(np.floor(_S_up / _Ly)) for _n_wrap in range(_n_lo, _n_up + 1): - _t = (_n_wrap * _Ly - _S_lo) / (_S_up - _S_lo) + _t = (_n_wrap * _Ly - _S_lo) / (_S_up - _S_lo) _x_wrap = _x_lo + _t * (_x_up - _x_lo) if _x_lo < _x_wrap < _x_up: - _ax.axvline(_x_wrap, color='limegreen', lw=1.2, ls=':', zorder=4, alpha=0.85) + _ax.axvline(_x_wrap, + color='limegreen', + lw=1.2, + ls=':', + zorder=4, + alpha=0.85) # ΔS/dy annotations for _ix, _dS in enumerate(_dS_cells): _color = 'red' if _dS >= _Ly else 'black' - _ax.text(_x_centers[_ix], _y_min - 0.06 * _Ly, f'{_dS/_dy:.1f}', - ha='center', va='top', fontsize=6, color=_color) - _ax.text(_x_min - 0.004 * _Lx, _y_min - 0.06 * _Ly, - r'$\Delta S/dy$:', ha='right', va='top', fontsize=7, color='black') + _ax.text(_x_centers[_ix], + _y_min - 0.06 * _Ly, + f'{_dS/_dy:.1f}', + ha='center', + va='top', + fontsize=6, + color=_color) + _ax.text(_x_min - 0.004 * _Lx, + _y_min - 0.06 * _Ly, + r'$\Delta S/dy$:', + ha='right', + va='top', + fontsize=7, + color='black') # Legend _legend_elems = [ - patches.Patch(fc='royalblue', alpha=0.4, label='donor cells'), - Line2D([0], [0], color='navy', lw=1.8, label=r'$y_{\rm tar,lo} - S(x)$'), - Line2D([0], [0], color='navy', lw=1.8, ls='--', label=r'$y_{\rm tar,up} - S(x)$'), - Line2D([0], [0], color='limegreen', lw=1.2, ls=':', label=r'wrap point: $S(x)=n\,L_y$'), + patches.Patch(fc='royalblue', alpha=0.4, label='donor cells'), + Line2D([0], [0], color='navy', lw=1.8, label=r'$y_{\rm tar,lo} - S(x)$'), + Line2D([0], [0], + color='navy', + lw=1.8, + ls='--', + label=r'$y_{\rm tar,up} - S(x)$'), + Line2D([0], [0], + color='limegreen', + lw=1.2, + ls=':', + label=r'wrap point: $S(x)=n\,L_y$'), ] - _ax.legend(handles=_legend_elems, loc='upper center', ncols=2, fontsize=9, - framealpha=0.92, frameon=True, bbox_to_anchor=(0.5, 1.0)) + _ax.legend(handles=_legend_elems, + loc='upper center', + ncols=2, + fontsize=9, + framealpha=0.92, + frameon=True, + bbox_to_anchor=(0.5, 1.0)) _ax.set_xlim(_x_min - 0.01 * _Lx, _x_max + 0.01 * _Lx) _ax.set_ylim(_y_min - 0.14 * _Ly, _y_max + 0.10 * _Ly) _ax.set_title( - f'twist-shift donor region for target $j=1$ (lowest y-cell)\n' - fr'$N_x={_Nx},~N_y={_Ny},~L_x={_Lx:.3e}$ m$,~L_y={_Ly:.3e}$ m', - fontsize=10, + f'twist-shift donor region for target $j=1$ (lowest y-cell)\n' + fr'$N_x={_Nx},~N_y={_Ny},~L_x={_Lx:.3e}$ m$,~L_y={_Ly:.3e}$ m', + fontsize=10, ) plt.tight_layout() @@ -344,7 +410,7 @@ def _field_plot(mo, pg, np, plt, result, species, quantity, comp_ui, boundary): mo.stop(True, mo.md("Run the simulation first.")) _ok, _sim_name, _sdir = result - _bnd = boundary.value # "lower" or "upper" + _bnd = boundary.value # "lower" or "upper" _comp = int(comp_ui.value) _field_file = f"{_sdir}/{_sim_name}-{species.value}_{quantity.value}_0.gkyl" @@ -353,13 +419,14 @@ def _field_plot(mo, pg, np, plt, result, species, quantity, comp_ui, boundary): try: _pg_data = pg.GData(_field_file) except Exception as _e: - mo.stop(True, mo.callout( - mo.md(f"**Cannot load file:**\n`{_field_file}`\n\n{_e}"), - kind="danger")) + mo.stop( + True, + mo.callout(mo.md(f"**Cannot load file:**\n`{_field_file}`\n\n{_e}"), + kind="danger")) # Determine poly_order and basis_type from file context _poly_order = _pg_data.ctx.get('poly_order', 1) - _ctx_basis = _pg_data.ctx.get('basis_type', 'serendipity') + _ctx_basis = _pg_data.ctx.get('basis_type', 'serendipity') _basis_type = 'ms' if _ctx_basis == 'serendipity' else _ctx_basis # Interpolate @@ -369,26 +436,30 @@ def _field_plot(mo, pg, np, plt, result, species, quantity, comp_ui, boundary): # _data_out: (Nx*p, Ny*p, Nz*p, 1) (with comp already selected → last dim = 1) # Slice at z boundary - _iz = 0 if _bnd == "lower" else -1 - _field = _data_out[:, :, _iz, 0] # shape (Nx_interp, Ny_interp) + _iz = 0 if _bnd == "lower" else -1 + _field = _data_out[:, :, _iz, 0] # shape (Nx_interp, Ny_interp) # Plot _fig, _ax = plt.subplots(figsize=(9, 5.5)) _ax.set_aspect('equal') # pcolormesh with nodal x/y grids: data must be (Ny, Nx) = transposed - _pcm = _ax.pcolormesh( - _x_out[0], _x_out[1], _field.T, - shading='flat', cmap='inferno') - plt.colorbar(_pcm, ax=_ax, label=f"{species.value}_{quantity.value} [comp {_comp}]") + _pcm = _ax.pcolormesh(_x_out[0], + _x_out[1], + _field.T, + shading='flat', + cmap='inferno') + plt.colorbar(_pcm, + ax=_ax, + label=f"{species.value}_{quantity.value} [comp {_comp}]") _ax.set_xlabel('x (m)', fontsize=12) _ax.set_ylabel('y (m)', fontsize=12) _bnd_label = "lower" if _bnd == "lower" else "upper" _ax.set_title( - f'species: {species.value}, quantity: {quantity.value}, ' - f'comp: {_comp}, z boundary: {_bnd_label}', - fontsize=10, + f'species: {species.value}, quantity: {quantity.value}, ' + f'comp: {_comp}, z boundary: {_bnd_label}', + fontsize=10, ) plt.tight_layout() @@ -411,30 +482,47 @@ def _ts_lib(mo, gkylsoft_ui, run_btn): _lib_gk = ctypes.CDLL(f"{_gkylsoft}/gkeyll/lib/libg0gyrokinetic.so") _c_vp = ctypes.c_void_p - _c_i = ctypes.c_int - _c_b = ctypes.c_bool + _c_i = ctypes.c_int + _c_b = ctypes.c_bool # gkyl_sub_range_init(rng*, bigrng*, sublower*, subupper*) - _lc.gkyl_sub_range_init.argtypes = [_c_vp, _c_vp, - ctypes.POINTER(_c_i), ctypes.POINTER(_c_i)] + _lc.gkyl_sub_range_init.argtypes = [ + _c_vp, _c_vp, ctypes.POINTER(_c_i), + ctypes.POINTER(_c_i) + ] _lc.gkyl_sub_range_init.restype = None - _lib_gk.gkyl_bc_twistshift_new.argtypes = [_c_i, _c_i, _c_i, _c_i, _c_i, _c_vp, - ctypes.POINTER(_c_i), _c_vp, _c_vp, _c_vp, _c_vp, _c_vp, _c_i, _c_b, ] - _lib_gk.gkyl_bc_twistshift_new.restype = _c_vp + _lib_gk.gkyl_bc_twistshift_new.argtypes = [ + _c_i, + _c_i, + _c_i, + _c_i, + _c_i, + _c_vp, + ctypes.POINTER(_c_i), + _c_vp, + _c_vp, + _c_vp, + _c_vp, + _c_vp, + _c_i, + _c_b, + ] + _lib_gk.gkyl_bc_twistshift_new.restype = _c_vp _lib_gk.gkyl_bc_twistshift_advance.argtypes = [_c_vp, _c_vp, _c_vp] - _lib_gk.gkyl_bc_twistshift_advance.restype = None + _lib_gk.gkyl_bc_twistshift_advance.restype = None _lib_gk.gkyl_bc_twistshift_release.argtypes = [_c_vp] - _lib_gk.gkyl_bc_twistshift_release.restype = None + _lib_gk.gkyl_bc_twistshift_release.restype = None ts_lib = (_lc, _lib_gk) return ctypes, ts_lib @app.cell -def _ts_plot(mo, ctypes, ts_lib, result, species, quantity, comp_ui, boundary, pg, np, plt): +def _ts_plot(mo, ctypes, ts_lib, result, species, quantity, comp_ui, boundary, + pg, np, plt): if result is None or not result[0]: mo.stop(True, mo.md("Run the simulation first.")) @@ -448,8 +536,10 @@ def _ts_plot(mo, ctypes, ts_lib, result, species, quantity, comp_ui, boundary, p try: _pg_data = pg.GData(_field_file) except Exception as _e: - mo.stop(True, mo.callout( - mo.md(f"**Cannot load field:**\n`{_field_file}`\n\n{_e}"), kind="danger")) + mo.stop( + True, + mo.callout(mo.md(f"**Cannot load field:**\n`{_field_file}`\n\n{_e}"), + kind="danger")) _poly_order = int(_pg_data.ctx.get('poly_order', 1)) _interior = np.ascontiguousarray(_pg_data.get_values(), dtype=np.float64) @@ -460,37 +550,41 @@ def _ts_plot(mo, ctypes, ts_lib, result, species, quantity, comp_ui, boundary, p _ext[1:-1, 1:-1, 1:-1, :] = _interior # Apply periodic BC along z before twist-shift: ghost ← opposite interior face - _ext[:, :, 0, :] = _ext[:, :, -2, :] # lower z ghost ← last interior z cell - _ext[:, :, -1, :] = _ext[:, :, 1, :] # upper z ghost ← first interior z cell + _ext[:, :, 0, :] = _ext[:, :, -2, :] # lower z ghost ← last interior z cell + _ext[:, :, -1, :] = _ext[:, :, 1, :] # upper z ghost ← first interior z cell # Load shift DG from file written by gyrokinetic_app_write_ts_shift _shift_file = f"{_sdir}/{_sim_name}-bc_z{_bnd}_twistshift.gkyl" try: _shift_pg = pg.GData(_shift_file) except Exception as _e: - mo.stop(True, mo.callout( - mo.md(f"**Cannot load shift file:**\n`{_shift_file}`\n\n{_e}"), kind="danger")) + mo.stop( + True, + mo.callout( + mo.md(f"**Cannot load shift file:**\n`{_shift_file}`\n\n{_e}"), + kind="danger")) _shift_poly_order = int(_shift_pg.ctx.get('poly_order', _poly_order)) _shift_vals = np.ascontiguousarray(_shift_pg.get_values(), dtype=np.float64) - _Nx_shift = int(np.prod(_shift_vals.shape[:-1])) + _Nx_shift = int(np.prod(_shift_vals.shape[:-1])) _shift_ncomp = _shift_vals.shape[-1] # Create gkyl objects (libg0core argtypes already set by GkeyllDGops) - _xn = _pg_data.get_grid() + _xn = _pg_data.get_grid() _lower = (ctypes.c_double * 3)(_xn[0][0], _xn[1][0], _xn[2][0]) _upper = (ctypes.c_double * 3)(_xn[0][-1], _xn[1][-1], _xn[2][-1]) - _cells = (ctypes.c_int * 3)(_Nx, _Ny, _Nz) - _ghost = (ctypes.c_int * 3)(1, 1, 1) + _cells = (ctypes.c_int * 3)(_Nx, _Ny, _Nz) + _ghost = (ctypes.c_int * 3)(1, 1, 1) - _basis_ptr = _lc.gkyl_cart_modal_serendip_new(ctypes.c_int(3), ctypes.c_int(_poly_order)) - _grid_ptr = _lc.gkyl_rect_grid_new(ctypes.c_int(3), _lower, _upper, _cells) + _basis_ptr = _lc.gkyl_cart_modal_serendip_new(ctypes.c_int(3), + ctypes.c_int(_poly_order)) + _grid_ptr = _lc.gkyl_rect_grid_new(ctypes.c_int(3), _lower, _upper, _cells) # local_ext: full extended range [0,Nx+1]x[0,Ny+1]x[0,Nz+1] matching _ext buffer strides _lo_ext = (ctypes.c_int * 3)(0, 0, 0) _up_ext = (ctypes.c_int * 3)(_Nx + 1, _Ny + 1, _Nz + 1) _local_ext = _lc.gkyl_range_new(ctypes.c_int(3), _lo_ext, _up_ext) - # bc_range: sub-range [1,Nx]x[1,Ny]x[0,Nz+1] — x/y interior only, z includes ghosts. + # bc_range: sub-range [1,Nx]x[1,Ny]x[0,Nz+1] -- x/y interior only, z includes ghosts. # Allocate via gkyl_range_new (opaque pointer), then let gkyl_sub_range_init overwrite # its fields in-place so it inherits _local_ext's linearizer (correct buffer strides). _lo_sub = (ctypes.c_int * 3)(1, 1, 0) @@ -500,19 +594,19 @@ def _ts_plot(mo, ctypes, ts_lib, result, species, quantity, comp_ui, boundary, p _lc.gkyl_range_release(_local_ext) # sub-range has copied the linearizer _f_arr = _lc.gkyl_array_new_from_buff( - 2, ctypes.c_size_t(_ncomp), ctypes.c_size_t((_Nx + 2) * (_Ny + 2) * (_Nz + 2)), - _ext.ctypes.data_as(ctypes.c_void_p)) + 2, ctypes.c_size_t(_ncomp), + ctypes.c_size_t((_Nx + 2) * (_Ny + 2) * (_Nz + 2)), + _ext.ctypes.data_as(ctypes.c_void_p)) _shift_arr = _lc.gkyl_array_new_from_buff( - 2, ctypes.c_size_t(_shift_ncomp), ctypes.c_size_t(_Nx_shift), - _shift_vals.ctypes.data_as(ctypes.c_void_p)) + 2, ctypes.c_size_t(_shift_ncomp), ctypes.c_size_t(_Nx_shift), + _shift_vals.ctypes.data_as(ctypes.c_void_p)) # Create twist-shift updater, apply in-place, then release everything _edge = 0 if _bnd == "lower" else 1 - _up = _lib_gk.gkyl_bc_twistshift_new( - 2, 1, 0, _edge, 3, - _bc_range, _ghost, _basis_ptr, _grid_ptr, - None, None, _shift_arr, _shift_poly_order, False) + _up = _lib_gk.gkyl_bc_twistshift_new(2, 1, 0, _edge, 3, _bc_range, _ghost, + _basis_ptr, _grid_ptr, None, None, + _shift_arr, _shift_poly_order, False) _lib_gk.gkyl_bc_twistshift_advance(_up, _f_arr, _f_arr) @@ -531,23 +625,29 @@ def _ts_plot(mo, ctypes, ts_lib, result, species, quantity, comp_ui, boundary, p _ghost_gdata.ctx = dict(_pg_data.ctx) _ghost_gdata.push([_xn[0], _xn[1], _xn[2][0:2]], _ghost_3d) - _ctx_basis = _pg_data.ctx.get('basis_type', 'serendipity') + _ctx_basis = _pg_data.ctx.get('basis_type', 'serendipity') _basis_type = 'ms' if _ctx_basis == 'serendipity' else _ctx_basis - _pg_interp = pg.GInterpModal(_ghost_gdata, _poly_order, _basis_type) + _pg_interp = pg.GInterpModal(_ghost_gdata, _poly_order, _basis_type) _x_interp, _data_interp = _pg_interp.interpolate(_comp) _ghost_field = _data_interp[:, :, 0, 0] # (Nx*p, Ny*p) _fig, _ax = plt.subplots(figsize=(9, 5.5)) _ax.set_aspect('equal') - _pcm = _ax.pcolormesh(_x_interp[0], _x_interp[1], _ghost_field.T, shading='flat', cmap='inferno') - plt.colorbar(_pcm, ax=_ax, label=f"{species.value}_{quantity.value} ghost [comp {_comp}]") + _pcm = _ax.pcolormesh(_x_interp[0], + _x_interp[1], + _ghost_field.T, + shading='flat', + cmap='inferno') + plt.colorbar(_pcm, + ax=_ax, + label=f"{species.value}_{quantity.value} ghost [comp {_comp}]") _ax.set_xlabel('x (m)', fontsize=12) _ax.set_ylabel('y (m)', fontsize=12) _bnd_label = "lower" if _bnd == "lower" else "upper" _ax.set_title( - f'Ghost after twist-shift — {species.value} {quantity.value}, ' - f'comp {_comp}, z {_bnd_label}', - fontsize=10, + f'Ghost after twist-shift -- {species.value} {quantity.value}, ' + f'comp {_comp}, z {_bnd_label}', + fontsize=10, ) plt.tight_layout() _out3 = mo.as_html(_fig) diff --git a/notebooks/twist_shift.mo.py b/notebooks/twist_shift.mo.py index 7200821c..70c260c9 100644 --- a/notebooks/twist_shift.mo.py +++ b/notebooks/twist_shift.mo.py @@ -3,135 +3,140 @@ __generated_with = "0.21.1" app = marimo.App(width="full") -# Cell: utility functions — defines the q-profile and low-pass filter used in the notebook. + +# Cell: utility functions -- defines the q-profile and low-pass filter used in the notebook. @app.cell def _(): - import numpy as np - - def filter_downsample(field, nint_x, nint_y, filter_type="lanczos", a=2, sigma=0.5, periodic_y=True): - """Downsample a 2-D field by integer factors with an antialiasing pre-filter. + import numpy as np + + def filter_downsample(field, + nint_x, + nint_y, + filter_type="lanczos", + a=2, + sigma=0.5, + periodic_y=True): + """Downsample a 2-D field by integer factors with an antialiasing pre-filter. Separable: filters along y (periodic) then x (non-periodic edge-padded). np.sinc is normalised: sinc(x) = sin(pi*x)/(pi*x), sinc(0) = 1. """ - if filter_type == "5point": - Nx_f, Ny_f = field.shape - tmp = np.zeros_like(field) - - # --- 1. filter along x (non-periodic, 2-point at boundaries) --- - if Nx_f >= 3: - tmp[1:-1, :] = (field[:-2, :] + field[1:-1, :] + field[2:, :]) / 3.0 - tmp[0, :] = (field[0, :] + field[1, :]) / 2.0 - tmp[-1, :] = (field[-2, :] + field[-1, :]) / 2.0 - else: - tmp[...] = field - - # --- 2. filter along y --- - out = np.zeros_like(tmp) - if Ny_f >= 3: - if periodic_y: - pad_y = np.concatenate([tmp[:, -1:], tmp, tmp[:, :1]], axis=1) - out = (pad_y[:, :-2] + pad_y[:, 1:-1] + pad_y[:, 2:]) / 3.0 - else: - out[:, 1:-1] = (tmp[:, :-2] + tmp[:, 1:-1] + tmp[:, 2:]) / 3.0 - out[:, 0] = (tmp[:, 0] + tmp[:, 1]) / 2.0 - out[:, -1] = (tmp[:, -2] + tmp[:, -1]) / 2.0 - else: - out[...] = tmp - - return out[::nint_x, ::nint_y] - - elif filter_type == "9point": - Nx_f, Ny_f = field.shape - mask = np.ones_like(field) - if periodic_y: - pf_y = np.concatenate([field[:, -1:], field, field[:, :1]], axis=1) - pm_y = np.concatenate([mask[:, -1:], mask, mask[:, :1]], axis=1) - else: - pf_y = np.pad(field, ((0, 0), (1, 1)), mode='constant') - pm_y = np.pad(mask, ((0, 0), (1, 1)), mode='constant') - - pad_field = np.pad(pf_y, ((1, 1), (0, 0)), mode='constant') - pad_mask = np.pad(pm_y, ((1, 1), (0, 0)), mode='constant') - - sum_field = ( - pad_field[:-2, :-2] + pad_field[:-2, 1:-1] + pad_field[:-2, 2:] + - pad_field[1:-1, :-2] + pad_field[1:-1, 1:-1] + pad_field[1:-1, 2:] + - pad_field[2:, :-2] + pad_field[2:, 1:-1] + pad_field[2:, 2:] - ) - sum_mask = ( - pad_mask[:-2, :-2] + pad_mask[:-2, 1:-1] + pad_mask[:-2, 2:] + - pad_mask[1:-1, :-2] + pad_mask[1:-1, 1:-1] + pad_mask[1:-1, 2:] + - pad_mask[2:, :-2] + pad_mask[2:, 1:-1] + pad_mask[2:, 2:] - ) - - out = sum_field / np.maximum(sum_mask, 1) - return out[::nint_x, ::nint_y] - - def make_kernel(n, filter_type, a, sigma): - if n == 1: - return np.ones(1) - if filter_type == "box": - # Uniform average of n cells. Simple but slow rolloff (~0.64 at Nyquist). - w = np.ones(n) - elif filter_type == "triangle": - # Convolution of two box filters: all-positive, squared-sinc frequency - # response — better Nyquist attenuation (~0.41) than box with no ringing. - box = np.ones(n) - w = np.convolve(box, box) # length 2n-1 - elif filter_type == "gaussian": - # All-positive, no ringing. sigma is the half-width in coarse-cell units. - half = int(np.ceil(3.0 * sigma * n)) - k = np.arange(-half, half + 1, dtype=float) - w = np.exp(-0.5 * (k / (sigma * n)) ** 2) - elif filter_type == "hann": - # Pure Hann window smoother (no sinc), all-positive: cosine rolloff over - # a*n fine cells. Smoother than Gaussian at the cost of a wider stencil. - half = a * n - k = np.arange(-half, half + 1, dtype=float) - w = 0.5 + 0.5 * np.cos(np.pi * k / half) - elif filter_type == "blackman": - # Pure Blackman window smoother (no sinc), all-positive: sharper rolloff - # than Hann, best stopband of the window-only family. - half = a * n - k = np.arange(-half, half + 1, dtype=float) - t = k / half - w = 0.42 + 0.5 * np.cos(np.pi * t) + 0.08 * np.cos(2 * np.pi * t) - else: - w = np.ones(n) - w = np.clip(w, 0, None) # ensure all-positive (safety guard) - w /= w.sum() - return w - - Nx_f, Ny_f = field.shape - - # --- filter along y (axis=1) --- - wy = make_kernel(nint_y, filter_type, a, sigma) - half_y = len(wy) // 2 + if filter_type == "5point": + Nx_f, Ny_f = field.shape + tmp = np.zeros_like(field) + + # --- 1. filter along x (non-periodic, 2-point at boundaries) --- + if Nx_f >= 3: + tmp[1:-1, :] = (field[:-2, :] + field[1:-1, :] + field[2:, :]) / 3.0 + tmp[0, :] = (field[0, :] + field[1, :]) / 2.0 + tmp[-1, :] = (field[-2, :] + field[-1, :]) / 2.0 + else: + tmp[...] = field + + # --- 2. filter along y --- + out = np.zeros_like(tmp) + if Ny_f >= 3: if periodic_y: - pad_y = np.concatenate([field[:, -half_y:], field, field[:, :half_y]], axis=1) + pad_y = np.concatenate([tmp[:, -1:], tmp, tmp[:, :1]], axis=1) + out = (pad_y[:, :-2] + pad_y[:, 1:-1] + pad_y[:, 2:]) / 3.0 else: - pad_y = np.pad(field, ((0, 0), (half_y, half_y)), mode='edge') - tmp = sum(wy[ki] * pad_y[:, ki:ki + Ny_f] for ki in range(len(wy))) + out[:, 1:-1] = (tmp[:, :-2] + tmp[:, 1:-1] + tmp[:, 2:]) / 3.0 + out[:, 0] = (tmp[:, 0] + tmp[:, 1]) / 2.0 + out[:, -1] = (tmp[:, -2] + tmp[:, -1]) / 2.0 + else: + out[...] = tmp + + return out[::nint_x, ::nint_y] + + elif filter_type == "9point": + Nx_f, Ny_f = field.shape + mask = np.ones_like(field) + if periodic_y: + pf_y = np.concatenate([field[:, -1:], field, field[:, :1]], axis=1) + pm_y = np.concatenate([mask[:, -1:], mask, mask[:, :1]], axis=1) + else: + pf_y = np.pad(field, ((0, 0), (1, 1)), mode='constant') + pm_y = np.pad(mask, ((0, 0), (1, 1)), mode='constant') + + pad_field = np.pad(pf_y, ((1, 1), (0, 0)), mode='constant') + pad_mask = np.pad(pm_y, ((1, 1), (0, 0)), mode='constant') + + sum_field = (pad_field[:-2, :-2] + pad_field[:-2, 1:-1] + + pad_field[:-2, 2:] + pad_field[1:-1, :-2] + + pad_field[1:-1, 1:-1] + pad_field[1:-1, 2:] + + pad_field[2:, :-2] + pad_field[2:, 1:-1] + pad_field[2:, 2:]) + sum_mask = (pad_mask[:-2, :-2] + pad_mask[:-2, 1:-1] + pad_mask[:-2, 2:] + + pad_mask[1:-1, :-2] + pad_mask[1:-1, 1:-1] + + pad_mask[1:-1, 2:] + pad_mask[2:, :-2] + pad_mask[2:, 1:-1] + + pad_mask[2:, 2:]) + + out = sum_field / np.maximum(sum_mask, 1) + return out[::nint_x, ::nint_y] + + def make_kernel(n, filter_type, a, sigma): + if n == 1: + return np.ones(1) + if filter_type == "box": + # Uniform average of n cells. Simple but slow rolloff (~0.64 at Nyquist). + w = np.ones(n) + elif filter_type == "triangle": + # Convolution of two box filters: all-positive, squared-sinc frequency + # response -- better Nyquist attenuation (~0.41) than box with no ringing. + box = np.ones(n) + w = np.convolve(box, box) # length 2n-1 + elif filter_type == "gaussian": + # All-positive, no ringing. sigma is the half-width in coarse-cell units. + half = int(np.ceil(3.0 * sigma * n)) + k = np.arange(-half, half + 1, dtype=float) + w = np.exp(-0.5 * (k / (sigma * n))**2) + elif filter_type == "hann": + # Pure Hann window smoother (no sinc), all-positive: cosine rolloff over + # a*n fine cells. Smoother than Gaussian at the cost of a wider stencil. + half = a * n + k = np.arange(-half, half + 1, dtype=float) + w = 0.5 + 0.5 * np.cos(np.pi * k / half) + elif filter_type == "blackman": + # Pure Blackman window smoother (no sinc), all-positive: sharper rolloff + # than Hann, best stopband of the window-only family. + half = a * n + k = np.arange(-half, half + 1, dtype=float) + t = k / half + w = 0.42 + 0.5 * np.cos(np.pi * t) + 0.08 * np.cos(2 * np.pi * t) + else: + w = np.ones(n) + w = np.clip(w, 0, None) # ensure all-positive (safety guard) + w /= w.sum() + return w + + Nx_f, Ny_f = field.shape + + # --- filter along y (axis=1) --- + wy = make_kernel(nint_y, filter_type, a, sigma) + half_y = len(wy) // 2 + if periodic_y: + pad_y = np.concatenate([field[:, -half_y:], field, field[:, :half_y]], + axis=1) + else: + pad_y = np.pad(field, ((0, 0), (half_y, half_y)), mode='edge') + tmp = sum(wy[ki] * pad_y[:, ki:ki + Ny_f] for ki in range(len(wy))) - # --- filter along x (axis=0, reflect-padded) --- - # reflect padding mirrors the signal at the boundary, so the filter - # averages real (mirrored) values at the edges rather than a constant, - # preventing amplitude preservation artifacts with edge padding. - wx = make_kernel(nint_x, filter_type, a, sigma) - half_x = len(wx) // 2 - pad_x = np.pad(tmp, ((half_x, half_x), (0, 0)), mode='reflect') - out = sum(wx[ki] * pad_x[ki:ki + Nx_f, :] for ki in range(len(wx))) + # --- filter along x (axis=0, reflect-padded) --- + # reflect padding mirrors the signal at the boundary, so the filter + # averages real (mirrored) values at the edges rather than a constant, + # preventing amplitude preservation artifacts with edge padding. + wx = make_kernel(nint_x, filter_type, a, sigma) + half_x = len(wx) // 2 + pad_x = np.pad(tmp, ((half_x, half_x), (0, 0)), mode='reflect') + out = sum(wx[ki] * pad_x[ki:ki + Nx_f, :] for ki in range(len(wx))) - return out[::nint_x, ::nint_y] + return out[::nint_x, ::nint_y] - def q_func(r, r0, q0, nshear): - return q0 * pow(r/r0, nshear) - - return filter_downsample, q_func, np + def q_func(r, r0, q0, nshear): + return q0 * pow(r / r0, nshear) + return filter_downsample, q_func, np -# Cell: introductory markdown — explains the twist-and-shift boundary condition and what the notebook does. +# Cell: introductory markdown -- explains the twist-and-shift boundary condition and what the notebook does. # @app.cell # def _(mo): # mo.md( @@ -141,7 +146,7 @@ def q_func(r, r0, q0, nshear): # This notebook lets you interactively explore twist-and-shift boundary conditions # used in flux-tube gyrokinetic simulations. -# Given a field $\phi(x, y)$ with a single $(k_x, k_y)$ wave, this notebook applies +# Given a field $\phi(x, y)$ with a single $(k_x, k_y)$ wave, this notebook applies # the twist-and-shift map directly in real space, using the periodic remapping # $$y \to y + 2\pi r_0/q_0 q(x).$$ @@ -153,374 +158,475 @@ def q_func(r, r0, q0, nshear): # return -# Cell: imports — loads marimo (reactive UI) and numpy. +# Cell: imports -- loads marimo (reactive UI) and numpy. @app.cell def _(): - import marimo as mo - return mo + import marimo as mo + return mo -# Cell: imports — loads matplotlib for all plotting. +# Cell: imports -- loads matplotlib for all plotting. @app.cell def _(): - import matplotlib.pyplot as plt - return plt + import matplotlib.pyplot as plt + return plt -# Cell: grid-size sliders — lets the user choose the number of grid points Nx and Ny. +# Cell: grid-size sliders -- lets the user choose the number of grid points Nx and Ny. @app.cell def _(mo): - Nx_slider = mo.ui.slider(value=32, start=4, stop=256, step=4, label="Nx (grid points in x)") - Ny_slider = mo.ui.slider(value=32, start=4, stop=256, step=4, label="Ny (grid points in y)") - return Nx_slider, Ny_slider - -# Cell: wave-parameter sliders — selects the field mode (plane wave or Gaussian packet) and + Nx_slider = mo.ui.slider(value=32, + start=4, + stop=256, + step=4, + label="Nx (grid points in x)") + Ny_slider = mo.ui.slider(value=32, + start=4, + stop=256, + step=4, + label="Ny (grid points in y)") + return Nx_slider, Ny_slider + + +# Cell: wave-parameter sliders -- selects the field mode (plane wave or Gaussian packet) and # the integer mode numbers kx_mode, ky_mode, plus the Gaussian envelope widths sigma_x/y. @app.cell def _(mo): - # --- Wave parameters --- - # field_mode: "single" is a pure cosine plane wave; "gaussian" is a - # localised wave-packet (Gaussian envelope × cosine carrier). - # kx_mode / ky_mode: integer mode numbers; the physical wave-numbers - # are kx = kx_mode * 2π/Lx and ky = ky_mode * 2π/Ly. - # sigma_x / sigma_y: half-widths of the Gaussian envelope (in box units); - # only used when field_mode == "gaussian". - field_mode = mo.ui.dropdown(options=["single", "gaussian"], value="single", label="Field mode") - kx_mode_slider = mo.ui.slider(value=2, start=0, stop=8, step=1, label="kx mode number (integer)") - ky_mode_slider = mo.ui.slider(value=1, start=0, stop=8, step=1, label="ky mode number (integer)") - sigma_x_slider = mo.ui.slider(value=0.1, start=0.05, stop=2.0, step=0.01, label="sigma_x (plane units)") - sigma_y_slider = mo.ui.slider(value=1.0, start=0.1, stop=6.0, step=0.01, label="sigma_y (plane units)") - add_mode2 = mo.ui.checkbox(value=True, label="Add 2nd mode") - kx2_mode_slider = mo.ui.slider(value=5, start=0, stop=8, step=1, label="kx2 mode number (integer)") - ky2_mode_slider = mo.ui.slider(value=5, start=0, stop=8, step=1, label="ky2 mode number (integer)") - return field_mode, kx_mode_slider, ky_mode_slider, sigma_x_slider, sigma_y_slider, add_mode2, kx2_mode_slider, ky2_mode_slider - -# Cell: magnetic-shear slider — sets the normalised shear ŝ that governs the twist-and-shift offset. + # --- Wave parameters --- + # field_mode: "single" is a pure cosine plane wave; "gaussian" is a + # localised wave-packet (Gaussian envelope × cosine carrier). + # kx_mode / ky_mode: integer mode numbers; the physical wave-numbers + # are kx = kx_mode * 2π/Lx and ky = ky_mode * 2π/Ly. + # sigma_x / sigma_y: half-widths of the Gaussian envelope (in box units); + # only used when field_mode == "gaussian". + field_mode = mo.ui.dropdown(options=["single", "gaussian"], + value="single", + label="Field mode") + kx_mode_slider = mo.ui.slider(value=2, + start=0, + stop=8, + step=1, + label="kx mode number (integer)") + ky_mode_slider = mo.ui.slider(value=1, + start=0, + stop=8, + step=1, + label="ky mode number (integer)") + sigma_x_slider = mo.ui.slider(value=0.1, + start=0.05, + stop=2.0, + step=0.01, + label="sigma_x (plane units)") + sigma_y_slider = mo.ui.slider(value=1.0, + start=0.1, + stop=6.0, + step=0.01, + label="sigma_y (plane units)") + add_mode2 = mo.ui.checkbox(value=True, label="Add 2nd mode") + kx2_mode_slider = mo.ui.slider(value=5, + start=0, + stop=8, + step=1, + label="kx2 mode number (integer)") + ky2_mode_slider = mo.ui.slider(value=5, + start=0, + stop=8, + step=1, + label="ky2 mode number (integer)") + return field_mode, kx_mode_slider, ky_mode_slider, sigma_x_slider, sigma_y_slider, add_mode2, kx2_mode_slider, ky2_mode_slider + + +# Cell: magnetic-shear slider -- sets the normalised shear ŝ that governs the twist-and-shift offset. @app.cell def _(mo): - shat_slider = mo.ui.slider(value=2.5, start=0.0, stop=5.0, step=0.5, label="ŝ (magnetic shear)") - q0_slider = mo.ui.slider(value=2.1, start=1.0, stop=5.0, step=0.1, label="q0 (safety factor)") - return shat_slider, q0_slider - -# Cell: visualisation-control widgets — colormap picker and oversampling factors for the shift. + shat_slider = mo.ui.slider(value=2.5, + start=0.0, + stop=5.0, + step=0.5, + label="ŝ (magnetic shear)") + q0_slider = mo.ui.slider(value=2.1, + start=1.0, + stop=5.0, + step=0.1, + label="q0 (safety factor)") + return shat_slider, q0_slider + + +# Cell: visualisation-control widgets -- colormap picker and oversampling factors for the shift. @app.cell def _(mo): - cmap_dropdown = mo.ui.dropdown(options=["seismic", "viridis", "twilight"], value="seismic", label="Colormap") - nint_x_slider = mo.ui.slider(value=4, start=1, stop=8, step=1, label="Oversampling x (nint_x)") - nint_y_slider = mo.ui.slider(value=4, start=1, stop=8, step=1, label="Oversampling y (nint_y)") - apply_filter = mo.ui.checkbox(value=True, label="Apply downsampling filter") - filter_dropdown = mo.ui.dropdown( - options=["gaussian", "triangle", "hann", "blackman", "box", "5point", "9point"], - value="gaussian", label="Downsampling filter", - ) - lanczos_a_slider = mo.ui.slider(value=2, start=1, stop=5, step=1, label="Lobe count a (lanczos/hann/blackman)") - gaussian_sigma_slider = mo.ui.slider(value=0.5, start=0.1, stop=3.0, step=0.1, label="Gaussian σ (in coarse cells)") - return cmap_dropdown, nint_x_slider, nint_y_slider, apply_filter, filter_dropdown, lanczos_a_slider, gaussian_sigma_slider - - -# Cell: core computation — builds the grid, constructs the test field phi(x,y), applies the + cmap_dropdown = mo.ui.dropdown(options=["seismic", "viridis", "twilight"], + value="seismic", + label="Colormap") + nint_x_slider = mo.ui.slider(value=4, + start=1, + stop=8, + step=1, + label="Oversampling x (nint_x)") + nint_y_slider = mo.ui.slider(value=4, + start=1, + stop=8, + step=1, + label="Oversampling y (nint_y)") + apply_filter = mo.ui.checkbox(value=True, label="Apply downsampling filter") + filter_dropdown = mo.ui.dropdown( + options=[ + "gaussian", "triangle", "hann", "blackman", "box", "5point", "9point" + ], + value="gaussian", + label="Downsampling filter", + ) + lanczos_a_slider = mo.ui.slider(value=2, + start=1, + stop=5, + step=1, + label="Lobe count a (lanczos/hann/blackman)") + gaussian_sigma_slider = mo.ui.slider(value=0.5, + start=0.1, + stop=3.0, + step=0.1, + label="Gaussian σ (in coarse cells)") + return cmap_dropdown, nint_x_slider, nint_y_slider, apply_filter, filter_dropdown, lanczos_a_slider, gaussian_sigma_slider + + +# Cell: core computation -- builds the grid, constructs the test field phi(x,y), applies the # twist-and-shift via spectral y-shifts, and optionally low-pass filters both fields. @app.cell -def _(field_mode,nint_x_slider,nint_y_slider,apply_filter,filter_dropdown,lanczos_a_slider,gaussian_sigma_slider, - Nx_slider,Ny_slider,kx_mode_slider,ky_mode_slider,np,shat_slider,q0_slider, - sigma_x_slider,sigma_y_slider,q_func,filter_downsample, - add_mode2,kx2_mode_slider,ky2_mode_slider): - - # --- Unpack slider values into plain variables --- - Nx = Nx_slider.value - Ny = Ny_slider.value - x0 = 2.0 # reference radial position (centre of the x-domain) - Lx = 1.0 # radial box size - A = 1.0 # target amplitude after normalisation - kx_mode = kx_mode_slider.value - ky_mode = ky_mode_slider.value - shat = shat_slider.value - q0 = q0_slider.value # safety factor at the reference surface x0 - Cy = x0/q0 - Ly = 2*np.pi*Cy - - # --- Build the 2-D spatial grid --- - dx = Lx / Nx - dy = Ly / Ny - x = np.linspace(x0, x0 + Lx - dx, Nx) - y = np.linspace(-Ly / 2, Ly / 2 - dy, Ny) - X, Y = np.meshgrid(x, y, indexing="ij") # shape (Nx, Ny) - - # Fundamental wave-numbers (smallest non-zero k on this grid) - dkx = 2.0 * np.pi / Lx - dky = 2.0 * np.pi / Ly - - # Physical wave-numbers of the chosen mode - kx_val = kx_mode * dkx - ky_val = ky_mode * dky - - # Local radial coordinate measured from the reference surface at x0. - # Used so that the shear shift is zero at x = x0. - x_local = x - x0 - - # --- safety-factor profile --- - q_profile = q_func(x, x0, q0, shat) - - # --- Construct the original 2-D test field phi(x, y) --- - if field_mode.value == "single": - # Pure plane wave: a single Fourier mode filling the whole domain. - phi = np.cos(kx_val * X + ky_val * Y) - else: - # Gaussian wave-packet: a localised envelope modulating the carrier. - # This is useful for visualising how a spatially-confined perturbation - # is remapped by the twist-and-shift. - x_center = x0 - y_center = 0.0 - sigma_x = sigma_x_slider.value - sigma_y = sigma_y_slider.value * Cy - - # Gaussian envelope centred at (x_center, y_center) - envelope = np.exp( - -0.5 * ((X - x_center) / sigma_x) ** 2 - -0.5 * ((Y - y_center) / sigma_y) ** 2 - ) - # Cosine carrier wave (the wave-like oscillation inside the packet) - carrier = np.cos(kx_val * (X - x_center) + ky_val * (Y - y_center)) - phi = envelope * carrier - - # Add optional second mode before normalisation. - if add_mode2.value: - kx2_val = kx2_mode_slider.value * dkx - ky2_val = ky2_mode_slider.value * dky - phi = phi + np.cos(kx2_val * X + ky2_val * Y) - - # Normalize to unit amplitude so the colour scale is the same for both modes. - phi_max = np.max(np.abs(phi)) - if phi_max > 0: - phi = A * phi / phi_max - - # Coarse-grid y-shift (kept for diagnostics) - delta_y = 2.0 * np.pi * Cy * q_profile # shape (Nx,) - - # --- 1. Upsample phi to fine grid (separable: periodic in y, linear in x) --- - nint_x = nint_x_slider.value - nint_y = nint_y_slider.value - Nx_fine = Nx * nint_x - Ny_fine = Ny * nint_y - dx_fine = Lx / Nx_fine - dy_fine = Ly / Ny_fine - x_fine = np.linspace(x0, x0 + Lx - dx_fine, Nx_fine) - y_fine = np.linspace(-Ly / 2, Ly / 2 - dy_fine, Ny_fine) - X_fine, Y_fine = np.meshgrid(x_fine, y_fine, indexing="ij") - - phi_fine_y = np.zeros((Nx, Ny_fine)) - for i in range(Nx): - phi_fine_y[i, :] = np.interp(y_fine, y, phi[i, :], period=Ly) - phi_fine = np.zeros((Nx_fine, Ny_fine)) - for j in range(Ny_fine): - phi_fine[:, j] = np.interp(x_fine, x, phi_fine_y[:, j]) - - # --- 2. Apply twist-and-shift on the fine grid --- - q_profile_fine = q_func(x_fine, x0, q0, shat) - delta_y_fine = 2.0 * np.pi * Cy * q_profile_fine - Y_shifted_fine = Y_fine + delta_y_fine[:, np.newaxis] - Y_shifted_wrapped_fine = (Y_shifted_fine + Ly / 2) % Ly - Ly / 2 - - phi_shifted_fine = np.zeros_like(phi_fine) - for i in range(Nx_fine): - phi_shifted_fine[i, :] = np.interp( - Y_shifted_wrapped_fine[i, :], y_fine, phi_fine[i, :], period=Ly - ) - - # --- 3. Downsample: optionally apply antialiasing filter before decimation --- - # When apply_filter is off, decimate directly (no convolution) so that the - # effect of aliasing is visible by toggling the checkbox. - a = lanczos_a_slider.value - sigma = gaussian_sigma_slider.value - if apply_filter.value: - phi_shifted = filter_downsample( - phi_shifted_fine, nint_x, nint_y, - filter_type=filter_dropdown.value, a=a, sigma=sigma, periodic_y=True - ) - else: - phi_shifted = phi_shifted_fine[::nint_x, ::nint_y] - - return (A, x0, Lx, Ly, Nx, Ny, X, Y, delta_y, dky, dx, dy, - field_mode, kx_val, ky_val, phi, phi_shifted, - X_fine, Y_fine, phi_fine, phi_shifted_fine, - q_profile, shat, sigma_x_slider, sigma_y_slider, x, x_local, y) - -# Cell: fine-grid field plots — pcolormesh of phi and phi_shifted on the oversampled mesh. +def _(field_mode, nint_x_slider, nint_y_slider, apply_filter, filter_dropdown, + lanczos_a_slider, gaussian_sigma_slider, Nx_slider, Ny_slider, + kx_mode_slider, ky_mode_slider, np, shat_slider, q0_slider, + sigma_x_slider, sigma_y_slider, q_func, filter_downsample, add_mode2, + kx2_mode_slider, ky2_mode_slider): + + # --- Unpack slider values into plain variables --- + Nx = Nx_slider.value + Ny = Ny_slider.value + x0 = 2.0 # reference radial position (centre of the x-domain) + Lx = 1.0 # radial box size + A = 1.0 # target amplitude after normalisation + kx_mode = kx_mode_slider.value + ky_mode = ky_mode_slider.value + shat = shat_slider.value + q0 = q0_slider.value # safety factor at the reference surface x0 + Cy = x0 / q0 + Ly = 2 * np.pi * Cy + + # --- Build the 2-D spatial grid --- + dx = Lx / Nx + dy = Ly / Ny + x = np.linspace(x0, x0 + Lx - dx, Nx) + y = np.linspace(-Ly / 2, Ly / 2 - dy, Ny) + X, Y = np.meshgrid(x, y, indexing="ij") # shape (Nx, Ny) + + # Fundamental wave-numbers (smallest non-zero k on this grid) + dkx = 2.0 * np.pi / Lx + dky = 2.0 * np.pi / Ly + + # Physical wave-numbers of the chosen mode + kx_val = kx_mode * dkx + ky_val = ky_mode * dky + + # Local radial coordinate measured from the reference surface at x0. + # Used so that the shear shift is zero at x = x0. + x_local = x - x0 + + # --- safety-factor profile --- + q_profile = q_func(x, x0, q0, shat) + + # --- Construct the original 2-D test field phi(x, y) --- + if field_mode.value == "single": + # Pure plane wave: a single Fourier mode filling the whole domain. + phi = np.cos(kx_val * X + ky_val * Y) + else: + # Gaussian wave-packet: a localised envelope modulating the carrier. + # This is useful for visualising how a spatially-confined perturbation + # is remapped by the twist-and-shift. + x_center = x0 + y_center = 0.0 + sigma_x = sigma_x_slider.value + sigma_y = sigma_y_slider.value * Cy + + # Gaussian envelope centred at (x_center, y_center) + envelope = np.exp(-0.5 * ((X - x_center) / sigma_x)**2 - 0.5 * + ((Y - y_center) / sigma_y)**2) + # Cosine carrier wave (the wave-like oscillation inside the packet) + carrier = np.cos(kx_val * (X - x_center) + ky_val * (Y - y_center)) + phi = envelope * carrier + + # Add optional second mode before normalisation. + if add_mode2.value: + kx2_val = kx2_mode_slider.value * dkx + ky2_val = ky2_mode_slider.value * dky + phi = phi + np.cos(kx2_val * X + ky2_val * Y) + + # Normalize to unit amplitude so the colour scale is the same for both modes. + phi_max = np.max(np.abs(phi)) + if phi_max > 0: + phi = A * phi / phi_max + + # Coarse-grid y-shift (kept for diagnostics) + delta_y = 2.0 * np.pi * Cy * q_profile # shape (Nx,) + + # --- 1. Upsample phi to fine grid (separable: periodic in y, linear in x) --- + nint_x = nint_x_slider.value + nint_y = nint_y_slider.value + Nx_fine = Nx * nint_x + Ny_fine = Ny * nint_y + dx_fine = Lx / Nx_fine + dy_fine = Ly / Ny_fine + x_fine = np.linspace(x0, x0 + Lx - dx_fine, Nx_fine) + y_fine = np.linspace(-Ly / 2, Ly / 2 - dy_fine, Ny_fine) + X_fine, Y_fine = np.meshgrid(x_fine, y_fine, indexing="ij") + + phi_fine_y = np.zeros((Nx, Ny_fine)) + for i in range(Nx): + phi_fine_y[i, :] = np.interp(y_fine, y, phi[i, :], period=Ly) + phi_fine = np.zeros((Nx_fine, Ny_fine)) + for j in range(Ny_fine): + phi_fine[:, j] = np.interp(x_fine, x, phi_fine_y[:, j]) + + # --- 2. Apply twist-and-shift on the fine grid --- + q_profile_fine = q_func(x_fine, x0, q0, shat) + delta_y_fine = 2.0 * np.pi * Cy * q_profile_fine + Y_shifted_fine = Y_fine + delta_y_fine[:, np.newaxis] + Y_shifted_wrapped_fine = (Y_shifted_fine + Ly / 2) % Ly - Ly / 2 + + phi_shifted_fine = np.zeros_like(phi_fine) + for i in range(Nx_fine): + phi_shifted_fine[i, :] = np.interp(Y_shifted_wrapped_fine[i, :], + y_fine, + phi_fine[i, :], + period=Ly) + + # --- 3. Downsample: optionally apply antialiasing filter before decimation --- + # When apply_filter is off, decimate directly (no convolution) so that the + # effect of aliasing is visible by toggling the checkbox. + a = lanczos_a_slider.value + sigma = gaussian_sigma_slider.value + if apply_filter.value: + phi_shifted = filter_downsample(phi_shifted_fine, + nint_x, + nint_y, + filter_type=filter_dropdown.value, + a=a, + sigma=sigma, + periodic_y=True) + else: + phi_shifted = phi_shifted_fine[::nint_x, ::nint_y] + + return (A, x0, Lx, Ly, Nx, Ny, X, Y, delta_y, dky, dx, dy, field_mode, kx_val, + ky_val, phi, phi_shifted, X_fine, Y_fine, phi_fine, phi_shifted_fine, + q_profile, shat, sigma_x_slider, sigma_y_slider, x, x_local, y) + + +# Cell: fine-grid field plots -- pcolormesh of phi and phi_shifted on the oversampled mesh. @app.cell def _(X_fine, Y_fine, cmap_dropdown, np, phi_fine, phi_shifted_fine, plt): - _cmap = cmap_dropdown.value - _vmax = np.max(np.abs(phi_fine)) - _panels = [ - ("φ fine grid", phi_fine), - ("φ fine grid shifted", phi_shifted_fine), - ] - _fig, _axes = plt.subplots(1, 2, figsize=(10, 4.5), squeeze=False) - _axes = _axes[0] + _cmap = cmap_dropdown.value + _vmax = np.max(np.abs(phi_fine)) + _panels = [ + ("φ fine grid", phi_fine), + ("φ fine grid shifted", phi_shifted_fine), + ] + _fig, _axes = plt.subplots(1, 2, figsize=(10, 4.5), squeeze=False) + _axes = _axes[0] + for _i, (_title, _data) in enumerate(_panels): + _im = _axes[_i].pcolormesh(X_fine, + Y_fine, + _data, + cmap=_cmap, + vmin=-_vmax, + vmax=_vmax, + shading="auto") + _axes[_i].set_title(_title) + _axes[_i].set_xlabel("x") + _axes[_i].set_ylabel("y") + plt.colorbar(_im, ax=_axes[_i], shrink=0.8) + plt.tight_layout() + fig_fine = plt.gcf() + plt.close(fig_fine) + return fig_fine, + + +# Cell: real-space field plots -- side-by-side pcolormesh of the original and twist-shifted phi(x,y). +@app.cell +def _(X, Y, cmap_dropdown, mo, np, phi, phi_shifted, plt): + # Side-by-side pcolormesh plots of the original and/or twist-shifted field. + # Both panels share the same colour scale (vmin=-vmax, vmax=vmax) so that + # amplitude differences are immediately visible. + _cmap = cmap_dropdown.value + _vmax = np.max(np.abs(phi)) + + # Build a list of (title, data) pairs for whichever panels are enabled + _panels = [] + _panels.append(("φ(x,y)", phi)) + _panels.append(("φ(x,y+S(x))", phi_shifted)) + + _n = len(_panels) + if _n > 0: + _ncols = min(_n, 4) + _nrows = int(np.ceil(_n / _ncols)) + _fig, _axes = plt.subplots(_nrows, + _ncols, + figsize=(5 * _ncols, 4.5 * _nrows), + squeeze=False) for _i, (_title, _data) in enumerate(_panels): - _im = _axes[_i].pcolormesh( - X_fine, Y_fine, _data, cmap=_cmap, vmin=-_vmax, vmax=_vmax, shading="auto" - ) - _axes[_i].set_title(_title) - _axes[_i].set_xlabel("x") - _axes[_i].set_ylabel("y") - plt.colorbar(_im, ax=_axes[_i], shrink=0.8) + _r, _c = divmod(_i, _ncols) + _ax = _axes[_r][_c] + _vm = _vmax if "err" not in _title.lower() else max( + np.max(np.abs(_data)), 1e-10) + _im = _ax.pcolormesh(X, + Y, + _data, + cmap=_cmap, + vmin=-_vm, + vmax=_vm, + shading="auto") + _ax.set_title(_title) + _ax.set_xlabel("x") + _ax.set_ylabel("y") + plt.colorbar(_im, ax=_ax, shrink=0.8) + # Hide unused axes + for _i in range(_n, _nrows * _ncols): + _r, _c = divmod(_i, _ncols) + _axes[_r][_c].set_visible(False) plt.tight_layout() - fig_fine = plt.gcf() - plt.close(fig_fine) - return fig_fine, + fig_realspace = plt.gcf() + plt.close(fig_realspace) + else: + fig_realspace = mo.md("*Enable at least one panel above.*") + return fig_realspace, -# Cell: real-space field plots — side-by-side pcolormesh of the original and twist-shifted phi(x,y). -@app.cell -def _( X, Y, cmap_dropdown, mo, np, phi, phi_shifted, plt): - # Side-by-side pcolormesh plots of the original and/or twist-shifted field. - # Both panels share the same colour scale (vmin=-vmax, vmax=vmax) so that - # amplitude differences are immediately visible. - _cmap = cmap_dropdown.value - _vmax = np.max(np.abs(phi)) - - # Build a list of (title, data) pairs for whichever panels are enabled - _panels = [] - _panels.append(("φ(x,y)", phi)) - _panels.append(("φ(x,y+S(x))", phi_shifted)) - - _n = len(_panels) - if _n > 0: - _ncols = min(_n, 4) - _nrows = int(np.ceil(_n / _ncols)) - _fig, _axes = plt.subplots( - _nrows, _ncols, figsize=(5 * _ncols, 4.5 * _nrows), squeeze=False - ) - for _i, (_title, _data) in enumerate(_panels): - _r, _c = divmod(_i, _ncols) - _ax = _axes[_r][_c] - _vm = _vmax if "err" not in _title.lower() else max(np.max(np.abs(_data)), 1e-10) - _im = _ax.pcolormesh( - X, Y, _data, cmap=_cmap, vmin=-_vm, vmax=_vm, shading="auto" - ) - _ax.set_title(_title) - _ax.set_xlabel("x") - _ax.set_ylabel("y") - plt.colorbar(_im, ax=_ax, shrink=0.8) - # Hide unused axes - for _i in range(_n, _nrows * _ncols): - _r, _c = divmod(_i, _ncols) - _axes[_r][_c].set_visible(False) - plt.tight_layout() - fig_realspace = plt.gcf() - plt.close(fig_realspace) - else: - fig_realspace = mo.md("*Enable at least one panel above.*") - return fig_realspace, - - -# Cell: 2-D Fourier spectra — pcolormesh of |φ̂(kx,ky)| for the original and shifted fields, +# Cell: 2-D Fourier spectra -- pcolormesh of |φ̂(kx,ky)| for the original and shifted fields, # showing how the twist-and-shift mixes Fourier modes. @app.cell def _(Lx, Ly, Nx, Ny, np, phi, phi_shifted, plt): - _phi_hat = np.fft.fftshift(np.fft.fft2(phi)) / (Nx * Ny) - _phi_rs_hat = np.fft.fftshift(np.fft.fft2(phi_shifted)) / (Nx * Ny) - - # Build centred kx/ky axes (fftshift reorders from [0,+k,...,-k] to [-k,...,0,...,+k]) - _kx_axis = np.fft.fftshift(np.fft.fftfreq(Nx, d=Lx / Nx)) * 2 * np.pi - _ky_axis = np.fft.fftshift(np.fft.fftfreq(Ny, d=Ly / Ny)) * 2 * np.pi - _KX, _KY = np.meshgrid(_kx_axis, _ky_axis, indexing="ij") - - _spectra = [ - ("Original |φ̂|", np.abs(_phi_hat)), - ("Shifted |φ̂|", np.abs(_phi_rs_hat)), - ] - - _fig, _axes = plt.subplots(1, 2, figsize=(10, 4.5), squeeze=False) - _axes = _axes[0] - for _i, (_title, _data) in enumerate(_spectra): - _vm = np.max(_data) or 1e-10 - _im = _axes[_i].pcolormesh( - _KX, _KY, _data, cmap="hot", vmin=0, vmax=_vm, shading="auto" - ) - _axes[_i].set_title(_title) - _axes[_i].set_xlabel("kx") - _axes[_i].set_ylabel("ky") - _axes[_i].set_ylim(0, np.max(_ky_axis)) - plt.colorbar(_im, ax=_axes[_i], shrink=0.8) - plt.tight_layout() - fig_fourier = plt.gcf() - plt.close(fig_fourier) - return fig_fourier, - - -# Cell: safety-factor and shear profiles — plots q(x) and the local normalised shear ŝ(x), + _phi_hat = np.fft.fftshift(np.fft.fft2(phi)) / (Nx * Ny) + _phi_rs_hat = np.fft.fftshift(np.fft.fft2(phi_shifted)) / (Nx * Ny) + + # Build centred kx/ky axes (fftshift reorders from [0,+k,...,-k] to [-k,...,0,...,+k]) + _kx_axis = np.fft.fftshift(np.fft.fftfreq(Nx, d=Lx / Nx)) * 2 * np.pi + _ky_axis = np.fft.fftshift(np.fft.fftfreq(Ny, d=Ly / Ny)) * 2 * np.pi + _KX, _KY = np.meshgrid(_kx_axis, _ky_axis, indexing="ij") + + _spectra = [ + ("Original |φ̂|", np.abs(_phi_hat)), + ("Shifted |φ̂|", np.abs(_phi_rs_hat)), + ] + + _fig, _axes = plt.subplots(1, 2, figsize=(10, 4.5), squeeze=False) + _axes = _axes[0] + for _i, (_title, _data) in enumerate(_spectra): + _vm = np.max(_data) or 1e-10 + _im = _axes[_i].pcolormesh(_KX, + _KY, + _data, + cmap="hot", + vmin=0, + vmax=_vm, + shading="auto") + _axes[_i].set_title(_title) + _axes[_i].set_xlabel("kx") + _axes[_i].set_ylabel("ky") + _axes[_i].set_ylim(0, np.max(_ky_axis)) + plt.colorbar(_im, ax=_axes[_i], shrink=0.8) + plt.tight_layout() + fig_fourier = plt.gcf() + plt.close(fig_fourier) + return fig_fourier, + + +# Cell: safety-factor and shear profiles -- plots q(x) and the local normalised shear ŝ(x), # letting the user verify how the linear q-model behaves across the radial box. @app.cell def _(np, plt, q_profile, shat, x_local): - _fig4, (_ax_q, _ax_s) = plt.subplots(1, 2, figsize=(10, 3.5)) - - _ax_q.plot(x_local, q_profile, "k-", lw=2) - _ax_q.set_xlabel("x - x0") - _ax_q.set_ylabel("q(x)") - _ax_q.set_title("Safety factor profile") - _ax_q.grid(True, alpha=0.3) - - # Shear profile: shat_local = (x_local/q) * dq/dx_local - # For linear q: dq/dx_local = shat (parameter) - # Guard against division by zero if q ever passes through zero. - _shat_local = x_local/q_profile * np.gradient(q_profile, x_local) - _ax_s.plot(x_local, _shat_local, "b-", lw=2, label="ŝ_local = (x/q) dq/dx") - _ax_s.axhline(shat, color="r", ls="--", label=f"ŝ parameter = {shat:.2f}") - _ax_s.set_xlabel("x - x0") - _ax_s.set_ylabel("ŝ(x)") - _ax_s.set_title("Local magnetic shear") - _ax_s.legend() - _ax_s.grid(True, alpha=0.3) - - plt.tight_layout() - fig_profiles = plt.gcf() - plt.close(fig_profiles) - return fig_profiles, - -# Cell: main layout — parameters on the left, plots on the right. + _fig4, (_ax_q, _ax_s) = plt.subplots(1, 2, figsize=(10, 3.5)) + + _ax_q.plot(x_local, q_profile, "k-", lw=2) + _ax_q.set_xlabel("x - x0") + _ax_q.set_ylabel("q(x)") + _ax_q.set_title("Safety factor profile") + _ax_q.grid(True, alpha=0.3) + + # Shear profile: shat_local = (x_local/q) * dq/dx_local + # For linear q: dq/dx_local = shat (parameter) + # Guard against division by zero if q ever passes through zero. + _shat_local = x_local / q_profile * np.gradient(q_profile, x_local) + _ax_s.plot(x_local, _shat_local, "b-", lw=2, label="ŝ_local = (x/q) dq/dx") + _ax_s.axhline(shat, color="r", ls="--", label=f"ŝ parameter = {shat:.2f}") + _ax_s.set_xlabel("x - x0") + _ax_s.set_ylabel("ŝ(x)") + _ax_s.set_title("Local magnetic shear") + _ax_s.legend() + _ax_s.grid(True, alpha=0.3) + + plt.tight_layout() + fig_profiles = plt.gcf() + plt.close(fig_profiles) + return fig_profiles, + + +# Cell: main layout -- parameters on the left, plots on the right. @app.cell def _( - Nx_slider, Ny_slider, - field_mode, kx_mode_slider, ky_mode_slider, sigma_x_slider, sigma_y_slider, - add_mode2, kx2_mode_slider, ky2_mode_slider, - shat_slider, q0_slider, - cmap_dropdown, nint_x_slider, nint_y_slider, apply_filter, filter_dropdown, lanczos_a_slider, gaussian_sigma_slider, - fig_realspace, fig_fourier, fig_profiles, fig_fine, + Nx_slider, + Ny_slider, + field_mode, + kx_mode_slider, + ky_mode_slider, + sigma_x_slider, + sigma_y_slider, + add_mode2, + kx2_mode_slider, + ky2_mode_slider, + shat_slider, + q0_slider, + cmap_dropdown, + nint_x_slider, + nint_y_slider, + apply_filter, + filter_dropdown, + lanczos_a_slider, + gaussian_sigma_slider, + fig_realspace, + fig_fourier, + fig_profiles, + fig_fine, mo, ): - _left = mo.vstack([ - mo.hstack([Nx_slider]), - mo.hstack([Ny_slider]), - mo.hstack([field_mode]), - mo.hstack([kx_mode_slider]), - mo.hstack([ky_mode_slider]), - mo.hstack([sigma_x_slider]), - mo.hstack([sigma_y_slider]), - mo.hstack([add_mode2]), - mo.hstack([kx2_mode_slider]), - mo.hstack([ky2_mode_slider]), - mo.hstack([q0_slider]), - mo.hstack([shat_slider]), - mo.hstack([cmap_dropdown]), - mo.hstack([nint_x_slider]), - mo.hstack([nint_y_slider]), - mo.hstack([apply_filter]), - mo.hstack([filter_dropdown]), - mo.hstack([lanczos_a_slider]), - mo.hstack([gaussian_sigma_slider]), - ]) - _right = mo.vstack([ - mo.md("## Profiles & Diagnostics"), - fig_profiles, - mo.md("## Real-Space Fields (fine grid)"), - fig_fine, - mo.md("## Real-Space Fields"), - fig_realspace, - mo.md("## Fourier Space"), - fig_fourier, - ]) - mo.Html(f""" + _left = mo.vstack([ + mo.hstack([Nx_slider]), + mo.hstack([Ny_slider]), + mo.hstack([field_mode]), + mo.hstack([kx_mode_slider]), + mo.hstack([ky_mode_slider]), + mo.hstack([sigma_x_slider]), + mo.hstack([sigma_y_slider]), + mo.hstack([add_mode2]), + mo.hstack([kx2_mode_slider]), + mo.hstack([ky2_mode_slider]), + mo.hstack([q0_slider]), + mo.hstack([shat_slider]), + mo.hstack([cmap_dropdown]), + mo.hstack([nint_x_slider]), + mo.hstack([nint_y_slider]), + mo.hstack([apply_filter]), + mo.hstack([filter_dropdown]), + mo.hstack([lanczos_a_slider]), + mo.hstack([gaussian_sigma_slider]), + ]) + _right = mo.vstack([ + mo.md("## Profiles & Diagnostics"), + fig_profiles, + mo.md("## Real-Space Fields (fine grid)"), + fig_fine, + mo.md("## Real-Space Fields"), + fig_realspace, + mo.md("## Fourier Space"), + fig_fourier, + ]) + mo.Html(f"""
{_left.text} @@ -533,4 +639,4 @@ def _( if __name__ == "__main__": - app.run() + app.run() diff --git a/proto/dash-install.sh b/proto/dash-install.sh index 15ecd457..1b40567c 100755 --- a/proto/dash-install.sh +++ b/proto/dash-install.sh @@ -1,4 +1,4 @@ pip install dash==0.34.0 # The core dash backend pip install dash-html-components==0.13.4 # HTML components pip install dash-core-components==0.41.0 # Supercharged components -pip install dash-table==3.1.11 # Interactive DataTable component (new!) \ No newline at end of file +pip install dash-table==3.1.11 # Interactive DataTable component (new!) diff --git a/pyproject.toml b/pyproject.toml index 78df6d23..3218be22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=61.0"] +requires = ["setuptools>=61.0", "numpy>=2.2.6"] build-backend = "setuptools.build_meta" [project] @@ -12,12 +12,15 @@ authors = [ description = "Python library and command-line tool for postprocessing (not only) Gkeyll data" dependencies = [ "click>=8.1.7", - "matplotlib>=3.7.0", - "msgpack>=1.0.3", - "numpy>=1.24.4,<2", - "scipy>=1.10.1", - "sympy>=1.12", - "tables>=3.8.0", + "matplotlib>=3.10.9", + "msgpack>=1.1.2", + "numpy>=2.2.6", + "scipy>=1.15.3", + "tables>=3.10.1", + "plotly>=6.7.0", + "kaleido>=1.3.0", + "pyvista>=0.48.4", + "imageio-ffmpeg>=0.6.0", ] readme = "README.md" license = {file = "LICENSE"} @@ -38,14 +41,34 @@ classifiers = [ "Programming Language :: Python", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Topic :: Scientific/Engineering", ] [project.optional-dependencies] -adios = ["adios2>=2.9.0,<2.10.0"] -test = ["pytest>=7.4.0"] +docs = [ + "sphinx>=8,<9", + "furo>=2024.8.6", + "myst-parser>=4,<5", +] +test = [ + "pre-commit==4.3.0", + "pytest>=9.0.3", + "pytest-cov>=7.0.0", + "pytest-timeout>=2.4.0", + "ruff>=0.12.11", + "build>=1.3.0", + "twine>=6.2.0", + # macOS CI only: isolates each test in its own forked process (see + # .github/workflows/test.yml's --forked flag) so a native-library abort + # deep in a compiled dependency (matplotlib/FreeType, VTK, ...) fails + # just that one test instead of taking down the whole session. + "pytest-forked>=1.6.0; sys_platform == 'darwin'", +] [project.urls] Documentation = "https://gkeyll.readthedocs.io/" @@ -53,10 +76,66 @@ Repository = "https://github.com/ammarhakim/postgkyl" "Bug Tracker" = "https://github.com/ammarhakim/postgkyl/issues" [project.scripts] -pgkyl = "postgkyl.pgkyl:cli" +pgkyl = "postgkyl.cli.app:cli" [tool.setuptools.dynamic] version = {attr = "postgkyl.__version__"} [tool.setuptools.packages.find] -where = ["src/"] \ No newline at end of file +where = ["src/"] + +[tool.setuptools.package-data] +"postgkyl" = ["py.typed"] +"postgkyl.render" = ["*.mplstyle", "*.js"] +# the compiled bridge (scripts/build_gpython.sh) + the extension source; the +# gpython shim itself lives in the gkeyll repo (GKEYLL_C_SHIM.md) +"postgkyl.gpython" = ["_gpython.so", "csrc/*.c"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = ["--strict-config", "--strict-markers", "-ra"] +xfail_strict = true +markers = [ + "compatibility: pure-Python contract run across every supported interpreter", + "native: requires the compiled Gkeyll bridge and native kernels", + "render: exercises a graphical or animation rendering backend", + "external_tool: invokes an external executable such as Chrome or ffmpeg", + "slow: integration test intentionally unsuitable for the fast test lane", +] +filterwarnings = [ + "error", + "ignore:Animation was deleted without rendering anything:UserWarning", + "ignore:FigureCanvasAgg is non-interactive, and thus cannot be shown:UserWarning", + "ignore:invalid value encountered in divide:RuntimeWarning", + # ev_ops' sqrt fallback path (see test_operations_evaluate.py) deliberately + # runs sqrt on raw modal coefficients, some of which are negative. + "ignore:invalid value encountered in sqrt:RuntimeWarning", + "ignore::scipy.optimize.OptimizeWarning", + # vtkmodules' own numpy_support helper, not our code; fires on every + # pyvista array upload under NumPy >= 2.5. + "ignore:Setting the shape on a NumPy array has been deprecated:DeprecationWarning", + # plotly's kaleido wrapper, not our code; fires on every static image + # export against kaleido's server-based renderer. + "ignore:The kopts argument is ignored if using a server\\.:UserWarning", + # CPython's own fork() safety check, not our code; pytest's capture + # machinery keeps background threads alive, so any test that forks a + # worker process (animate's --nproc) trips it regardless of correctness. + "ignore:This process \\(pid=\\d+\\) is multi-threaded, use of fork\\(\\) may lead to deadlocks in the child\\.:DeprecationWarning", +] + +[tool.coverage.run] +branch = true +source = ["postgkyl"] + +[tool.coverage.report] +fail_under = 99 +precision = 2 +show_missing = true + +[tool.ruff] +target-version = "py310" +line-length = 80 +extend-exclude = ["gkeyll"] + +[tool.ruff.lint] +select = ["F"] diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 65c21d18..00000000 --- a/requirements.txt +++ /dev/null @@ -1,9 +0,0 @@ -click>=8.1.7 -conda-forge::adios2==2.9.2 -matplotlib>=3.7.0 -msgpack-python>=1.0.3 -numpy>=1.24.4 -pytables>=3.8.0 -pytest>=7.4.0 -scipy>=1.10.1 -sympy>=1.12 \ No newline at end of file diff --git a/scripts/build_docs.py b/scripts/build_docs.py new file mode 100644 index 00000000..934610c7 --- /dev/null +++ b/scripts/build_docs.py @@ -0,0 +1,317 @@ +"""Prepare Postgkyl's Sphinx subtree from this checkout and execute its gallery. + +Run with Python 3.12 after installing ``.[docs]``. The output is disposable; +an ownership marker prevents deleting a directory that this script did not make. +""" + +from __future__ import annotations + +import argparse +import inspect +import json +from pathlib import Path +import re +import runpy +import shlex +import shutil +import subprocess +import sys +import textwrap +from types import ModuleType +import zipfile + + +def public_functions(module: ModuleType, seen_modules: set, path: str): + """Yield public function paths through declared module exports.""" + if module in seen_modules: + return + seen_modules.add(module) + names = getattr(module, "__all__", sorted(vars(module))) + for name in names: + if name.startswith("_"): + continue + value = getattr(module, name) + if inspect.isfunction(value) and value.__module__.startswith("postgkyl"): + yield f"{path}.{name}", value + elif module.__name__.startswith("postgkyl.diagnostics") and isinstance( + value, + ModuleType) and value.__name__.startswith("postgkyl.diagnostics"): + yield from public_functions(value, seen_modules, f"{path}.{name}") + + +def write_reference(output: Path) -> dict: + """Render the real public functions and compiled Click command inventory.""" + import click + import postgkyl as pg + from postgkyl.cli.app import cli, COMMAND_SECTIONS, MODELS + + reference = output / "reference" + reference.mkdir() + api = [ + "Python API\n==========\n", + "Signatures and descriptions below come from the installed source.\n" + ] + seen_functions = {} + api_pages = {} + command_names = {} + for model in MODELS: + command_names.setdefault(model.canonical, []).append(model.name) + public_modules = { + value: f"postgkyl.{name}" + for name in pg.__all__ + if isinstance(value := getattr(pg, name), ModuleType) + } + roots = [("Core", pg)] + [(f"{name} diagnostics", getattr( + pg.diagnostics, name)) for name in pg.diagnostics.__all__] + seen_modules = set() + for title, module in roots: + api.extend( + [f"{title}\n{'-' * len(title)}\n", ".. toctree::\n :maxdepth: 1\n"]) + aliases = [] + module_path = public_modules.get(module, module.__name__) + for path, function in public_functions(module, seen_modules, module_path): + if function in seen_functions: + aliases.append( + f"``{path}`` is an alias of :func:`{seen_functions[function]}`.\n") + continue + seen_functions[function] = path + page = "api-" + path.replace(".", "-") + api_pages[path] = page + (reference / f"{page}.rst" + ).write_text(f"{path}\n{'=' * len(path)}\n\n.. autofunction:: {path}\n") + with (reference / f"{page}.rst").open("a") as page_file: + for name in command_names.get(function, ()): + page_file.write(f"\n:doc:`CLI: {name} `\n") + api.append(f" {page}") + api.extend(["", *aliases]) + api.append( + "Fluent datasets\n---------------\n\n.. toctree::\n :maxdepth: 1\n") + for name in ("GData", "GDataGroup"): + cls = getattr(pg, name) + class_page = f"api-{name}" + api.append(f" {class_page}") + members = [ + f"{name}\n{'=' * len(name)}\n", f".. autoclass:: postgkyl.{name}\n\n", + ".. toctree::\n :maxdepth: 1\n" + ] + aliases = [] + for member, value in inspect.getmembers(cls): + if member.startswith("_") or not (callable(value) + or isinstance(value, property)): + continue + path = f"postgkyl.{name}.{member}" + if value in seen_functions: + aliases.append(f"``{path}`` uses :func:`{seen_functions[value]}`.\n") + continue + page = "api-" + path.replace(".", "-") + kind = "autoattribute" if isinstance(value, property) else "automethod" + (reference / f"{page}.rst" + ).write_text(f"{path}\n{'=' * len(path)}\n\n.. {kind}:: {path}\n") + api_pages[path] = page + members.append(f" {page}") + (reference / f"{class_page}.rst").write_text("\n".join( + [*members, "", *aliases])) + api.append("") + (reference / "api.rst").write_text("\n".join(api)) + + commands = [] + contents = [ + "Command reference\n=================\n", + "This inventory and every help block are generated from the " + "same Click commands used by ``pgkyl --help``.\n" + ] + for section, names in COMMAND_SECTIONS.items(): + contents.extend([ + section, "-" * len(section), "", ".. toctree::", " :maxdepth: 1", "" + ]) + for name in names: + command = cli.commands[name] + with click.Context(command, info_name=f"pgkyl {name}", + terminal_width=88) as context: + help_text = command.get_help(context) + title = f"pgkyl {name}" + page = (f"{title}\n{'=' * len(title)}\n\n.. code-block:: text\n\n" + + textwrap.indent(help_text, " ") + "\n") + model = next(model for model in MODELS if model.name == name) + python_path = seen_functions.get(model.canonical) + if python_path is not None: + page += f"\n:doc:`Python API <{api_pages[python_path]}>`\n" + (reference / f"cli-{name}.rst").write_text(page) + contents.append(f" cli-{name}") + commands.append(name) + contents.append("") + (reference / "cli.rst").write_text("\n".join(contents)) + quantities = "\n".join(f"* ``{name}``" + for name in pg.gk.available_quantities()) + (reference / "quantities.rst").write_text( + "Gyrokinetic quantities\n======================\n\n" + "Names below come from ``pg.gk.available_quantities()``. " + f"See :func:`{seen_functions[pg.gk.load_quantity]}` for loading and " + "physical parameters.\n\n" + quantities + "\n") + return { + "api": list(seen_functions.values()), + "api_pages": api_pages, + "commands": commands + } + + +def write_figure_pairs(root: Path, output: Path, comparison: dict) -> None: + """Publish the commands that passed parity, beside the actual two outputs.""" + commands = json.loads((root / "examples/figure_commands.json").read_text()) + pairs = output / "_pairs" + pairs.mkdir() + interactive = output / "interactive" + interactive.mkdir() + report = [ + "Python and CLI figure comparisons\n=================================\n", + "Every result below was checked during this documentation build. " + "Raster comparisons decode pixels and animation timings; Plotly " + "comparisons check trace data, layout, and configuration. They " + "do not compare browser-specific rasterization.\n", + ".. list-table::\n :header-rows: 1\n\n" + " * - Output\n - Comparison\n" + ] + for script, outputs in commands.items(): + snippet = [ + "**Python script**\n", ".. code-block:: bash\n\n" + f" python examples/scripts/{script}\n", + f".. literalinclude:: _inputs/examples/scripts/{script}\n" + " :language: python\n", "**Equivalent CLI commands**\n", + "Run from the extracted example bundle (or repository root). " + "Create an output directory first:\n", + ".. code-block:: bash\n\n mkdir -p output\n" + ] + for name, command in outputs.items(): + arguments = [ + arg.replace("{output}", "output") for arg in shlex.split(command) + ] + # Wrap only between tokens; shell quoting remains correct for math labels. + lines = [] + current = "" + for arg in arguments: + token = shlex.quote(arg) + if current and len(current) + len(token) > 90: + lines.append(current + " \\") + current = " " + token + else: + current += (" " if current else "") + token + lines.append(current) + snippet.append(".. code-block:: bash\n\n" + + textwrap.indent("\n".join(lines), " ") + "\n") + snippet.append(f"``{name}``: {comparison[name]}.\n") + report.append(f" * - ``{name}``\n - {comparison[name]}\n") + if name.endswith(".html"): + for label, source in (("Python", output / "figures" / name), + ("CLI", output / "figures/cli" / name)): + target = f"{label.lower()}-{name}" + shutil.copyfile(source, interactive / target) + snippet.append( + f"{label}: :download:`open interactive figure `.\n" + ) + snippet.append( + ".. raw:: html\n\n" + f' \n') + else: + snippet.append( + ".. list-table:: Python and CLI outputs\n :widths: 50 50\n\n" + f" * - .. image:: figures/{name}\n" + f" :alt: Python output for {name}\n" + " :width: 100%\n" + f" - .. image:: figures/cli/{name}\n" + f" :alt: CLI output for {name}\n" + " :width: 100%\n") + (pairs / f"{Path(script).stem}.inc").write_text("\n".join(snippet)) + (output / "interface-equivalence.rst").write_text("\n".join(report)) + + +def prepare(root: Path, output: Path) -> None: + """Execute this checkout's examples and stage their sources and outputs.""" + import postgkyl as pg + from postgkyl import gpython + + if not Path(pg.__file__).resolve().is_relative_to(root / "src"): + raise RuntimeError( + "Install this checkout with pip install -e '.[docs]' first") + gpython.require() + marker = output / ".postgkyl-docs" + if output.exists(): + if not marker.is_file(): + raise ValueError( + f"Refusing to replace unowned output directory: {output}") + shutil.rmtree(output) + shutil.copytree(root / "docs" / "source", output) + marker.touch() + inputs = output / "_inputs" + shutil.copytree(root / "examples", + inputs / "examples", + ignore=shutil.ignore_patterns("output", "__pycache__")) + + subprocess.run([sys.executable, + str(root / "tests/generate_test_data.py")], + cwd=root, + check=True, + stdout=subprocess.DEVNULL) + figures = output / "figures" + figures.mkdir() + gallery = runpy.run_path(str(root / "examples/compare_interfaces.py")) + comparison = gallery["run_gallery"](root, figures) + write_figure_pairs(root, output, comparison) + shutil.copytree(root / "docs/_ext", output / "_ext") + + readme = (root / "README.md").read_text() + installation = readme.split("## Installation\n", + 1)[1].split("## Documentation\n", 1)[0] + notes = readme.split("### Additional installation notes\n", + 1)[1].split("## Developing for Postgkyl\n", 1)[0] + installation = "# Installation\n" + installation + "## Additional notes\n" + notes + installation = re.sub(r"^#(#{2,} )", r"\1", installation, flags=re.MULTILINE) + for name in ("environment.yml", "pyproject.toml"): + installation = installation.replace( + f"]({name})", + f"](https://github.com/ammarhakim/postgkyl/blob/main/{name})") + (output / "installation.md").write_text(installation) + + inventory = write_reference(output) + revision = subprocess.check_output( + ["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() + inventory.update(version=pg.__version__, revision=revision) + (output / + "build-info.json").write_text(json.dumps(inventory, indent=2) + "\n") + (output / "provenance.rst").write_text( + "Documentation source\n====================\n\n" + "The hosted documentation is rebuilt from Postgkyl's ``main`` branch. " + "Local and pull-request previews use their working checkout.\n\n" + f"This build used version ``{pg.__version__}``, commit ``{revision}``. " + "Uncommitted local changes, if present, are included in local previews.\n\n" + "`Edit the guides in Postgkyl " + "`_. " + "API descriptions are edited in the implementing Python functions.\n") + + # One download preserves the paths used by both the scripts and CLI tutorial. + downloads = output / "downloads" + downloads.mkdir() + with zipfile.ZipFile(downloads / "postgkyl-examples.zip", + "w", + compression=zipfile.ZIP_DEFLATED) as bundle: + files = list((inputs / "examples").rglob("*")) + for path in sorted(files): + if path.is_file(): + bundle.write(path, path.relative_to(inputs)) + for path in sorted((root / "tests/test_data").rglob("*.gkyl")): + bundle.write(path, path.relative_to(root)) + bundle.write(root / "tests/generate_test_data.py", + "tests/generate_test_data.py") + bundle.write(output / "build-info.json", "build-info.json") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, default=Path("build/docs/source")) + args = parser.parse_args() + prepare(Path(__file__).resolve().parents[1], args.output.resolve()) + + +if __name__ == "__main__": + main() diff --git a/scripts/build_gkeyll.sh b/scripts/build_gkeyll.sh new file mode 100755 index 00000000..1c4b9e34 --- /dev/null +++ b/scripts/build_gkeyll.sh @@ -0,0 +1,92 @@ +#!/bin/sh +# Fetches (if needed) and builds the vendored Gkeyll `core` app as +# libg0core.so, for the gpython/ layer to bind against. Invoked automatically +# by `pip install`/`pip install -e` via setup.py, and safe to re-run by hand. +# +# gkeyll/ is a plain, detached clone pinned by scripts/gkeyll-revision (zero +# external deps: no MPI/CUDA/SuperLU/Lua, LAPACK replaced by the bundled +# lapack-lite). Only core/ is needed to build libg0core.so, so moments/, +# vlasov/, gyrokinetic/, and pkpm/ (~200MB combined) are excluded via +# sparse-checkout and are never fetched, not merely deleted after the fact. +set -e + +REPO_URL="https://github.com/ammarhakim/gkeyll.git" +SPARSE_DIRS="core gkeyll install-deps machines" + +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +ROOT_DIR=$(CDPATH= cd -- "${SCRIPT_DIR}/.." && pwd) +GKEYLL_DIR="${ROOT_DIR}/gkeyll" +REVISION_FILE="${SCRIPT_DIR}/gkeyll-revision" + +if [ ! -f "${REVISION_FILE}" ]; then + echo "error: pinned Gkeyll revision file is missing: ${REVISION_FILE}" >&2 + exit 1 +fi +IFS= read -r GKEYLL_REVISION < "${REVISION_FILE}" +case "${GKEYLL_REVISION}" in + *[!0-9a-f]*|'') + echo "error: ${REVISION_FILE} must contain one lowercase commit SHA" >&2 + exit 1 + ;; +esac +if [ "${#GKEYLL_REVISION}" -ne 40 ]; then + echo "error: ${REVISION_FILE} must contain a full 40-character commit SHA" >&2 + exit 1 +fi + +if [ ! -e "${GKEYLL_DIR}/.git" ]; then + echo "# gkeyll/ not present -- fetching pinned ${GKEYLL_REVISION} (core-only, sparse + blobless)" + rmdir "${GKEYLL_DIR}" 2>/dev/null || true + mkdir "${GKEYLL_DIR}" + git -C "${GKEYLL_DIR}" init + git -C "${GKEYLL_DIR}" remote add origin "${REPO_URL}" + git -C "${GKEYLL_DIR}" sparse-checkout init --cone + git -C "${GKEYLL_DIR}" sparse-checkout set ${SPARSE_DIRS} + git -C "${GKEYLL_DIR}" fetch --depth 1 --filter=blob:none origin "${GKEYLL_REVISION}" +else + echo "# gkeyll/ already present -- ensuring sparse-checkout excludes heavy apps" + (cd "${GKEYLL_DIR}" && git sparse-checkout init --cone >/dev/null 2>&1 || true + git -C "${GKEYLL_DIR}" sparse-checkout set ${SPARSE_DIRS}) + if ! git -C "${GKEYLL_DIR}" cat-file -e "${GKEYLL_REVISION}^{commit}" 2>/dev/null; then + git -C "${GKEYLL_DIR}" fetch --depth 1 --filter=blob:none origin "${GKEYLL_REVISION}" + fi +fi + +# A dirty producer tree makes the native artifact's source unknowable even +# when HEAD is pinned. Refuse it instead of recording misleading build info. +if ! git -C "${GKEYLL_DIR}" diff --quiet || \ + ! git -C "${GKEYLL_DIR}" diff --cached --quiet; then + echo "error: ${GKEYLL_DIR} has tracked modifications; cannot build the pinned Gkeyll source" >&2 + exit 1 +fi +git -C "${GKEYLL_DIR}" checkout --detach "${GKEYLL_REVISION}" +ACTUAL_REVISION=$(git -C "${GKEYLL_DIR}" rev-parse HEAD) +if [ "${ACTUAL_REVISION}" != "${GKEYLL_REVISION}" ]; then + echo "error: expected Gkeyll ${GKEYLL_REVISION}, checked out ${ACTUAL_REVISION}" >&2 + exit 1 +fi +echo "# Using pinned Gkeyll revision ${GKEYLL_REVISION}" + +CC="${CC:-cc}" +echo "# Configuring gkeyll core (CC=${CC}, lapack-lite, app=core)" +(cd "${GKEYLL_DIR}" && ./configure "CC=${CC}" --use-lapack-lite=yes --app=core) + +ARCH_FLAGS="${ARCH_FLAGS:-}" +export ARCH_FLAGS + +echo "# Building libg0core.so (ARCH_FLAGS=${ARCH_FLAGS:-})" +(cd "${GKEYLL_DIR}" && make core "ARCH_FLAGS=${ARCH_FLAGS}" \ + -j"$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)") + +SO_PATH="${GKEYLL_DIR}/build/core/libg0core.so" +if [ ! -f "${SO_PATH}" ]; then + echo "error: expected ${SO_PATH} after build, but it is missing" >&2 + exit 1 +fi +echo "# Built ${SO_PATH}" + +# Build the _gpython extension against gkyl_gpython.h + libg0core.so. The +# gpython shim itself (core/zero/gpython.c) was just compiled INTO +# libg0core.so above -- that step is the compile-time contract check +# (GKEYLL_C_SHIM.md). +sh "${SCRIPT_DIR}/build_gpython.sh" diff --git a/scripts/build_gpython.sh b/scripts/build_gpython.sh new file mode 100755 index 00000000..dc5a1547 --- /dev/null +++ b/scripts/build_gpython.sh @@ -0,0 +1,126 @@ +#!/bin/sh +# Builds the _gpython CPython extension into src/postgkyl/gpython/_gpython.so +# (GKEYLL_C_SHIM.md). The gpython shim itself lives in the gkeyll repo +# (core/zero/gkyl_gpython.h + core/zero/gpython.c) and is compiled INTO +# libg0core.so by gkeyll's own build -- that compile step is the contract +# check: any core API drift fails there, at the producer. This script only +# compiles the extension against gkyl_gpython.h (opaque handles + scalars) and +# links the shim symbols from libg0core.so; a stale header/library pairing +# is caught at import by the GPYTHON_API_VERSION handshake. +# +# Requires a built gkeyll/build/core/libg0core.so (scripts/build_gkeyll.sh, +# which invokes this script as its final step). Safe to re-run by hand. +set -e + +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +ROOT_DIR=$(CDPATH= cd -- "${SCRIPT_DIR}/.." && pwd) +GKEYLL_DIR="${ROOT_DIR}/gkeyll" +LIB_DIR="${GKEYLL_DIR}/build/core" +CSRC_DIR="${ROOT_DIR}/src/postgkyl/gpython/csrc" +OUT="${ROOT_DIR}/src/postgkyl/gpython/_gpython.so" +BUNDLED_LIB="${ROOT_DIR}/src/postgkyl/gpython/libg0core.so" + +if [ ! -f "${LIB_DIR}/libg0core.so" ]; then + echo "error: ${LIB_DIR}/libg0core.so not found; run scripts/build_gkeyll.sh first" >&2 + exit 1 +fi +if [ ! -f "${GKEYLL_DIR}/core/zero/gkyl_gpython.h" ]; then + echo "error: gkeyll/core/zero/gkyl_gpython.h not found; this gkeyll tree lacks the gpython shim" >&2 + exit 1 +fi + +PYTHON="${PYTHON:-python3}" +PY_INCLUDES=$("${PYTHON}" -c "import sysconfig; print(sysconfig.get_path('include'))") +NUMPY_INCLUDE=$("${PYTHON}" -c "import numpy; print(numpy.get_include())") + +# _gpythonmodule.c targets NPY_2_2_API_VERSION (pyproject.toml's numpy>=2.2.6 +# floor) as a best-effort backstop, but a build-time/run-time NumPy version +# skew has been reproduced to crash the extension outright (segfault / heap +# corruption inside Gkeyll's C code) rather than fail cleanly -- the pin +# alone does not make a mismatched build safe. Under pip's default build +# isolation, this ${PYTHON} is a throwaway environment that resolves +# build-system.requires' numpy independently of whatever NumPy ends up +# installed for running postgkyl, so a stale/mismatched NumPy here (e.g. an +# ambient `python3` found ahead of the intended venv on PATH) would otherwise +# silently bake a broken extension instead of failing at build time. Always +# build with `--no-build-isolation` against the NumPy you're actually going +# to run with (see README.md). +NUMPY_OK=$("${PYTHON}" -c " +import sys +import numpy +major, minor = (int(p) for p in numpy.__version__.split('.')[:2]) +sys.stdout.write('yes' if (major, minor) >= (2, 2) else 'no') +") +if [ "${NUMPY_OK}" != "yes" ]; then + NUMPY_VERSION=$("${PYTHON}" -c "import numpy; print(numpy.__version__)") + echo "error: building _gpython against NumPy ${NUMPY_VERSION} (via ${PYTHON}), but postgkyl requires numpy>=2.2.6 (pyproject.toml)." >&2 + echo " This usually means '${PYTHON}' resolved to a different environment than the one postgkyl is being installed into." >&2 + echo " Reinstall with: pip install -e . --no-build-isolation" >&2 + exit 1 +fi + +CC="${CC:-cc}" + +# Keep the extension and its sole non-system shared library together. A +# relative loader path then works from a wheel, virtualenv, or relocated source +# checkout without referring back to this build tree. +cp "${LIB_DIR}/libg0core.so" "${BUNDLED_LIB}" + +# CPython extension modules must leave the Py* symbols unresolved at link +# time; the interpreter provides them at import. Linux's -shared does this +# by default, macOS needs -undefined dynamic_lookup (same flag setuptools +# passes on Darwin). +EXT_LDFLAGS="" +if [ "$(uname -s)" = "Darwin" ]; then + EXT_LDFLAGS="-Wl,-undefined,dynamic_lookup" + EXT_RPATH="@loader_path" + # Gkeyll names the Mach-O library .so for consistency across platforms. + # Give the bundled copy a relocatable install name before linking to it. + install_name_tool -id "@rpath/libg0core.so" "${BUNDLED_LIB}" +else + EXT_RPATH='$ORIGIN' +fi + +echo "# Building _gpython extension (CC=${CC}) -> ${OUT}" +"${CC}" -O2 -g -fPIC -shared \ + "${CSRC_DIR}/_gpythonmodule.c" \ + -I "${GKEYLL_DIR}/core/zero" \ + -I "${PY_INCLUDES}" \ + -I "${NUMPY_INCLUDE}" \ + -L "$(dirname -- "${BUNDLED_LIB}")" -lg0core -Wl,-rpath,"${EXT_RPATH}" \ + ${EXT_LDFLAGS} \ + -o "${OUT}" +echo "# Built ${OUT}" + +# Record what this build was made from -- gkeyll/ is a build-time-only clone +# (.gitignore'd), so its commit is otherwise unrecoverable once installed. +# Read by postgkyl._version for `pgkyl --version`. +BUILD_INFO="${ROOT_DIR}/src/postgkyl/gpython/_build_info.py" +_git_log_field() { # _git_log_field + git -C "$1" log -1 --format="$2" 2>/dev/null || echo unknown +} +GKEYLL_COMMIT=$(_git_log_field "${GKEYLL_DIR}" "%H") +GKEYLL_COMMIT_DATE=$(_git_log_field "${GKEYLL_DIR}" "%cI") +GKEYLL_BRANCH=pinned +POSTGKYL_BUILD_COMMIT=$(_git_log_field "${ROOT_DIR}" "%H") +BUILD_DATE=$(date -u +%Y-%m-%dT%H:%M:%SZ) +BUILD_ARCH_FLAGS="${ARCH_FLAGS:-}" + +cat > "${BUILD_INFO}" <&2 + exit 2 +fi + +PYTHON="${PYTHON:-python3}" +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +ROOT_DIR=$(CDPATH= cd -- "${SCRIPT_DIR}/.." && pwd) +SMOKE_FIELD="${ROOT_DIR}/tests/test_data/rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl" +case "$1" in + /*) WHEEL=$1 ;; + *) WHEEL=$(pwd)/$1 ;; +esac +if [ ! -f "${WHEEL}" ]; then + echo "error: wheel not found: ${WHEEL}" >&2 + exit 1 +fi + +SMOKE_DIR=$(mktemp -d "${TMPDIR:-/tmp}/postgkyl-wheel-smoke.XXXXXX") +trap 'rm -rf "${SMOKE_DIR}"' EXIT HUP INT TERM +if [ "${POSTGKYL_SMOKE_NO_DEPS:-0}" = "1" ]; then + # Useful for an offline developer check when the invoking interpreter + # already has postgkyl's dependencies. CI/release checks should leave + # this unset and exercise normal dependency resolution. + "${PYTHON}" -m venv --system-site-packages "${SMOKE_DIR}/venv" + "${SMOKE_DIR}/venv/bin/python" -m pip install --no-deps "${WHEEL}" +else + "${PYTHON}" -m venv "${SMOKE_DIR}/venv" + "${SMOKE_DIR}/venv/bin/python" -m pip install "${WHEEL}" +fi + +cd "${SMOKE_DIR}" +POSTGKYL_SMOKE_FIELD="${SMOKE_FIELD}" \ + "${SMOKE_DIR}/venv/bin/python" - <<'PY' +import os +from pathlib import Path + +import postgkyl +from postgkyl import gpython + +assert gpython.available(), "the wheel's compiled Gkeyll bridge did not load" +extension = gpython.lib_path() +assert extension is not None +core = extension.with_name("libg0core.so") +assert core.is_file(), f"wheel does not contain {core}" +assert "site-packages" in str(Path(postgkyl.__file__).resolve()) +assert gpython.require().api_version() == gpython.require().GPYTHON_API_VERSION +field = Path(os.environ["POSTGKYL_SMOKE_FIELD"]) +modal = postgkyl.load(field) +point_values = modal.interpolate() +assert modal.backend == "gkyl" +assert point_values.backend == "numpy" +print(f"loaded {extension}") +print(f"loaded bundled {core}") +PY +"${SMOKE_DIR}/venv/bin/pgkyl" --version +"${SMOKE_DIR}/venv/bin/python" -m pip check diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..843c07cd --- /dev/null +++ b/setup.py @@ -0,0 +1,104 @@ +import os +import subprocess +import sys +from pathlib import Path + +from setuptools import setup +from setuptools.dist import Distribution +from setuptools.command.build_py import build_py +from setuptools.command.develop import develop + +ROOT_DIR = Path(__file__).parent +BUILD_SCRIPT = ROOT_DIR / "scripts" / "build_gkeyll.sh" +BUNDLED_LIB = ROOT_DIR / "src" / "postgkyl" / "gpython" / "libg0core.so" +SKIP_BUILD_ENV = "POSTGKYL_SKIP_GKEYLL_BUILD" + + +def _skip_gkeyll_build(): + """Read the explicit pure-Python CI escape hatch. + + Only ``0`` and ``1`` are accepted so a misspelled value cannot silently + turn a required native packaging job into a pure-Python build. + """ + value = os.environ.get(SKIP_BUILD_ENV, "0") + if value not in {"0", "1"}: + raise RuntimeError(f"{SKIP_BUILD_ENV} must be 0 or 1, got {value!r}") + return value == "1" + + +def _build_gkeyll(): + if _skip_gkeyll_build(): + print(f"# Skipping Gkeyll build ({SKIP_BUILD_ENV}=1)") + return False + # build_gpython.sh (invoked as build_gkeyll.sh's final step) resolves its + # NumPy via `${PYTHON:-python3}` -- a bare PATH lookup. Left unset, that + # can silently resolve to a *different* Python installation than the one + # actually running this install (macOS commonly has several: system + # /usr/bin/python3, Homebrew, the actions/setup-python framework build). + # Such a mismatch passes build_gpython.sh's own `numpy>=2.2` floor check + # yet still isn't ABI-identical to the NumPy installed at runtime, which + # has been reproduced to crash the compiled _gpython extension outright + # (segfault / heap corruption) at some later, unrelated call rather than + # fail cleanly at import -- see scripts/build_gpython.sh and README.md. + # Pinning PYTHON to sys.executable here removes the ambiguity entirely: + # the extension is always built against exactly the NumPy this same + # interpreter will import at runtime. + env = os.environ.copy() + env["PYTHON"] = sys.executable + subprocess.run(["sh", str(BUILD_SCRIPT)], check=True, cwd=ROOT_DIR, env=env) + return True + + +class BuildPyWithGkeyll(build_py): + + def run(self): + built_native = _build_gkeyll() + super().run() + # PEP 660 sets editable_mode and deliberately makes build_py.run() a + # no-op: the editable wheel points at the source tree, where the build + # script has already placed both native artifacts. Trying to copy into + # build_lib here assumes a directory Setuptools intentionally did not + # create and makes ``pip install -e`` fail after a successful compile. + if getattr(self, "editable_mode", False): + return + destination = Path(self.build_lib) / "postgkyl" / "gpython" + if not built_native: + # A reused build directory may contain output from an earlier native + # build. A skip-build lane must remain genuinely pure Python. + for name in ("_build_info.py", "_gpython.so", "libg0core.so"): + (destination / name).unlink(missing_ok=True) + return + # libg0core is not a Python extension as far as setuptools knows, so copy + # it explicitly next to _gpython.so. The extension's relative loader path + # is deliberately resolved against this exact directory. + if not BUNDLED_LIB.is_file(): + raise FileNotFoundError( + f"native build did not produce bundled library: {BUNDLED_LIB}") + destination.mkdir(parents=True, exist_ok=True) + self.copy_file(str(BUNDLED_LIB), str(destination / BUNDLED_LIB.name)) + + +class DevelopWithGkeyll(develop): + + def run(self): + _build_gkeyll() + super().run() + + +class BinaryDistribution(Distribution): + """Mark wheels as interpreter/platform-specific native artifacts. + + ``_gpython.so`` is built by the Gkeyll-aware shell script rather than a + setuptools ``Extension``, so setuptools cannot infer this itself. Without + this declaration it emits an incorrect ``py3-none-any`` wheel. + """ + + def has_ext_modules(self): + return not _skip_gkeyll_build() + + +setup(cmdclass={ + "build_py": BuildPyWithGkeyll, + "develop": DevelopWithGkeyll, +}, + distclass=BinaryDistribution) diff --git a/src/postgkyl/__init__.py b/src/postgkyl/__init__.py index 4151d082..e68778a4 100644 --- a/src/postgkyl/__init__.py +++ b/src/postgkyl/__init__.py @@ -1,23 +1,99 @@ -""" -# Postgkyl +"""postgkyl -- a small, layered post-processing library for Gkeyll data. -Postgkyl is both Python library and command-line tool designed to provide unified access -to Gkeyll data together with a broad variety of analytical and visualization tools. -""" +Public surface (the facade). The golden script:: + + import postgkyl as pg + pg.load('elc_M0_0.gkyl').interpolate().select(z0=0.0).plot() + +The facade is **pure re-export** -- every public name is defined in the layer that +owns it and simply gathered here: -__version__ = "1.7.5" + load, GData, GDataGroup <- api/ (fluent surface) + collect, evaluate, relchange, sort <- api/ (module-level multi-dataset + verbs with no single ``self``) + plot, animate, plotly, <- render/ (canonical render callables) + plotly_animate, pyvista + group_blocks <- gdatastate/ (multiblock family partition) + info <- operations/ (the info verb, one-or-many) + integrate <- operations/ (grid integral, via Gkeyll) + interpolate, select <- operations/ (functional verb spellings) + gk_rz <- operations/gyrokinetics/ (domain operation) + represent, apply <- operations/ (value_form verbs) + available_evaluate_operators <- operations/ (``evaluate``'s RPN token vocabulary) + save <- io/ (file output) + gk <- diagnostics/ (equation namespace) + version_report <- _version.py (``pgkyl --version``'s + commit/build-info report) + +Every computational fluent ``GData`` method delegates to one of these +``operations`` functions, so ``pg.select(a, z0=0.0)`` and +``a.select(z0=0.0)`` are the same call -- the functional and fluent spellings +can never drift apart. ``GData.load(...)`` is the one lifecycle method: it +loads a literal file into an existing object and returns that same object for +chaining. The rest of the +domain-independent ``operations`` verb inventory (``fft``, ``magsq``, ``mask``, +``val2coord``, ``extract_input``, ``fit``, ``differentiate``, ``integrate``, +``map``, plus ``grid`` -- see ``api/gdata.py`` for why ``grid`` has no fluent +spelling) is reachable as a ``GData`` fluent method and via +``postgkyl.operations.``; this facade does not additionally promote each one to +a bare top-level name (one home per verb-vocabulary fact, not three). + +Architecture (strict, cycle-free DAG; see REFACTOR_GKEYLL_FFI.md):: + + floor gpython/ compiled _gpython extension -> libg0core.so (the only foreign code) + leaves numerics/ (pure NumPy; imports nothing internal) + engine dg/ interpolation bridge + modal ops -> gpython + leaves io/ readers (C-native first) -> gpython + container gdatastate/ GDataState {gkyl|numpy} backend + seam operations/ one verb each + backend render/ Matplotlib, Plotly, PyVista + fluent api/ GData(GDataState) + operators <- above operations + facade __init__ re-exports only +""" -# import submodules -from postgkyl import data -from postgkyl import utils -from postgkyl import tools -from postgkyl import output +from postgkyl.gdata import GData, load, GDataGroup, collect, evaluate, relchange, sort +from postgkyl.operations import ( + apply, + available_evaluate_operators, + average, + differentiate, + eval_at_coord_proj, + extract_input, + fft, + fit, + grid, + growth, + info, + integrate, + interpolate, + local_poly, + magsq, + map, + mask, + print, + represent, + select, + val2coord, +) +from postgkyl.operations.gyrokinetics import gk_fluxsurf, gk_rz +from postgkyl.render import animate, plot, plotly, plotly_animate, pyvista +from postgkyl.gdatastate import group_blocks +from postgkyl.cli_spec import hidden +from postgkyl.io import save +from postgkyl.diagnostics import gk +from postgkyl._version import version_report -# import selected classes to the root -from postgkyl.data.gdata import GData -from postgkyl.data.dg import GInterpNodal -from postgkyl.data.dg import GInterpModal +__version__ = "2.0.0" -# link the command line executable to the system -from postgkyl import pgkyl +hidden("collection helper is a Python API, not a pipeline command")( + group_blocks) +__all__ = [ + "GData", "load", "GDataGroup", "plot", "group_blocks", "info", "print", + "integrate", "interpolate", "local_poly", "select", "average", + "eval_at_coord_proj", "fft", "magsq", "mask", "grid", "val2coord", + "extract_input", "fit", "growth", "differentiate", "map", "represent", + "apply", "gk_rz", "gk_fluxsurf", "save", "collect", "evaluate", "relchange", + "animate", "plotly_animate", "sort", "available_evaluate_operators", + "plotly", "pyvista", "gk", "__version__", "version_report" +] diff --git a/src/postgkyl/_gkylsoft_path.py b/src/postgkyl/_gkylsoft_path.py deleted file mode 100644 index d13f1162..00000000 --- a/src/postgkyl/_gkylsoft_path.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Default gkylsoft path, baked in at install time or edited post-install. - -Ways to specify gkylsoft path (from highest to lowest priority): - 1. Manually specified (alt_gkylsoft_dir argument). - 2. GKYLSOFT_DIR environment variable. - 3. GKYLSOFT_DIR=... in the config file (default ~/.postgkyl/gkylsoft_path, - overridden by the POSTGKYL_CONFIG environment variable). - 4. GKYLSOFT_DIR below (set at install time, or edit this file directly). -""" - -GKYLSOFT_DIR = "" - -def default_config_path() -> str: - """Return the config file path, respecting POSTGKYL_CONFIG if set.""" - import os - return os.environ.get("POSTGKYL_CONFIG", - os.path.expanduser("~/.postgkyl/gkylsoft_path")) - -def resolve_gkylsoft_path(alt_gkylsoft_dir: str | None = None) -> str | None: - """Return the gkylsoft directory path, or None if not configured.""" - import os - - if alt_gkylsoft_dir: - return alt_gkylsoft_dir - - env = os.environ.get("GKYLSOFT_DIR") - if env: - return env - - cfg = default_config_path() - if os.path.isfile(cfg): - text = open(cfg).read().strip() - if text: - return text.split("=")[1] - - return GKYLSOFT_DIR if GKYLSOFT_DIR else None diff --git a/src/postgkyl/_version.py b/src/postgkyl/_version.py new file mode 100644 index 00000000..108ee1f3 --- /dev/null +++ b/src/postgkyl/_version.py @@ -0,0 +1,97 @@ +"""Debugging statistics behind ``pgkyl --version`` -- not part of the public +computing API (only ``cli/app.py``'s ``--version`` flag reads this). + +Reports the postgkyl commit this checkout is at, the vendored Gkeyll commit +it was built against (via ``gpython.build_info()``, generated at build time +by scripts/build_gpython.sh since gkeyll/ is a build-time-only clone), and +interpreter/platform/dependency versions -- everything a bug report needs +without asking the user to gather it by hand. +""" + +from __future__ import annotations + +import importlib.metadata +import pathlib +import platform +import subprocess + +from postgkyl import gpython +from postgkyl.cli_spec import hidden + +_DEPENDENCIES = ("numpy", "scipy", "click", "matplotlib", "msgpack", "plotly", + "pyvista") + + +def _git(repo_dir: pathlib.Path, *args: str) -> str | None: + if not (repo_dir / ".git").is_dir(): + return None + try: + result = subprocess.run(["git", "-C", str(repo_dir), *args], + capture_output=True, + text=True, + timeout=5, + check=True) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() or None + + +def _postgkyl_commit() -> str: + # this file is src/postgkyl/_version.py -- the repo root is two levels up + repo_dir = pathlib.Path(__file__).resolve().parents[2] + commit = _git(repo_dir, "rev-parse", "--short=12", "HEAD") + if commit is None: + build = gpython.build_info() + baked = build["postgkyl_build_commit"] if build else None + if baked and baked != "unknown": + return f"{baked[:12]} (baked at build time, not a git checkout)" + return "unknown (not a git checkout)" + dirty = _git(repo_dir, "status", "--porcelain", "--untracked-files=no") + return f"{commit}{'-dirty' if dirty else ''}" + + +def _gkeyll_info() -> str: + build = gpython.build_info() + if build is None: + return "not built (no compiled Gkeyll bridge -- see scripts/build_gkeyll.sh)" + return (f"{build['gkeyll_commit'][:12]} ({build['gkeyll_branch']}, " + f"committed {build['gkeyll_commit_date']})") + + +def _dependency_versions() -> str: + versions = [] + for name in _DEPENDENCIES: + try: + versions.append(f"{name} {importlib.metadata.version(name)}") + except importlib.metadata.PackageNotFoundError: + continue + return ", ".join(versions) + + +def version_report(version: str) -> str: + """Build the full ``pgkyl --version`` report. + + Args: + version: Installed postgkyl version string. + + Returns: + A multiline environment and build report suitable for bug reports. + """ + build = gpython.build_info() + bridge = "available" if gpython.available() else "unavailable" + if build is not None: + arch = build["build_arch_flags"] or "compiler default" + bridge += f" (built {build['build_date']}, CC={build['build_cc']}, ARCH_FLAGS={arch})" + return "\n".join([ + f"pgkyl, version {version}", + f"postgkyl commit: {_postgkyl_commit()}", + f"Gkeyll: {_gkeyll_info()}", + f"gpython bridge: {bridge}", + f"Python: {platform.python_implementation()} {platform.python_version()}", + f"Platform: {platform.platform()}", + f"Dependencies: {_dependency_versions()}", + ]) + + +hidden("version reporting is handled by the manual --version front end")( + version_report) diff --git a/src/postgkyl/cli/__init__.py b/src/postgkyl/cli/__init__.py new file mode 100644 index 00000000..34e06b22 --- /dev/null +++ b/src/postgkyl/cli/__init__.py @@ -0,0 +1,5 @@ +"""CLI layer -- a chained Click pipeline over the public API (top SURFACES layer).""" + +from .app import cli + +__all__ = ["cli"] diff --git a/src/postgkyl/cli/app.py b/src/postgkyl/cli/app.py new file mode 100644 index 00000000..dbc6262f --- /dev/null +++ b/src/postgkyl/cli/app.py @@ -0,0 +1,131 @@ +"""``pgkyl`` command-line entry point -- a chained pipeline on pure Click. + +The chained syntax mirrors the fluent script API 1:1:: + + pg.load('f.gkyl').interpolate().select(z0=0).plot() # script + pgkyl f.gkyl interpolate select --z0 0 plot # CLI + +Chaining and callback-before-dispatch are native to ``click.Group(chain=True)``, +so the only custom code is a small :class:`PgkylGroup.get_command` override for +command-name abbreviation and treating a bare filename as an implicit ``load``. +Every subcommand is compiled from a public API callable at import time. This +module owns only chaining, command aliases, and bare-file dispatch; it +contains no per-command option or execution definitions. +""" + +from __future__ import annotations + +from glob import glob +from types import MappingProxyType + +import click + +from postgkyl import __version__, version_report +from postgkyl.cli.compiler import ( + build_click_command, + compile_public_surface, + group_by_section, +) +from postgkyl.cli.discovery import discover_public_surface +from postgkyl.cli.state import DataSpace + +# Compilation validates the complete discovered surface before registration. +# These aliases add spellings only; they never replace a generated command or +# alter its options. +MODELS = compile_public_surface(discover_public_surface()) +COMMANDS = tuple(build_click_command(model) for model in MODELS) +COMMAND_SECTIONS = group_by_section(MODELS) +COMMAND_ALIASES = MappingProxyType({"pl": "plot", "ev": "evaluate"}) + + +class PgkylGroup(click.Group): + """Click's chained group with spelling-only command aliases.""" + + def get_command(self, ctx, name): + cmd = super().get_command(ctx, name) + if cmd is not None: + return cmd + if name in COMMAND_ALIASES: + target = COMMAND_ALIASES[name] + command = super().get_command(ctx, target) + if command is not None: + return command + matches = [c for c in self.list_commands(ctx) if c.startswith(name)] + if len(matches) == 1: + return super().get_command(ctx, matches[0]) + if matches: + ctx.fail(f"Ambiguous command '{name}': {', '.join(sorted(matches))}") + return None + + def resolve_command(self, ctx, args): + """Expand a bare file pattern to the canonical ``load --file_name`` form.""" + if args: + token = args[0] + exact = click.Group.get_command(self, ctx, token) + alias = COMMAND_ALIASES.get(token) + if exact is None and alias is None and glob(token): + args[:1] = ["load", "--file_name", token] + return super().resolve_command(ctx, args) + + def format_commands(self, ctx, formatter) -> None: + """Group ``pgkyl --help``'s command listing under section headers. + + Presentation only (see ``commands/__init__.py``'s ``COMMAND_SECTIONS`` + and "14-cli.md"'s "Help output organization"): every command stays a + flat, chainable top-level ``click.Command`` resolved exactly as before; + only how they are *printed* changes, mirroring how ``git``/``docker`` + group their subcommand help. + """ + for section, names in COMMAND_SECTIONS.items(): + rows = [] + for name in names: + cmd = self.get_command(ctx, name) + if cmd is None: + continue + rows.append((name, cmd.get_short_help_str(limit=formatter.width - 6))) + if rows: + with formatter.section(section): + formatter.write_dl(rows) + + +def _print_version(ctx, param, value) -> None: + if not value or ctx.resilient_parsing: + return + click.echo(version_report(__version__)) + ctx.exit() + + +@click.group(cls=PgkylGroup, + chain=True, + context_settings=dict(help_option_names=["-h", "--help"])) +@click.option("--version", + is_flag=True, + expose_value=False, + is_eager=True, + callback=_print_version, + help="Show version, commit, Gkeyll build info, and exit.") +@click.pass_context +def cli(ctx) -> None: + """Postprocessing and plotting tool for Gkeyll data. + + Datasets are loaded, processed and plotted by chaining commands, e.g.:: + + pgkyl file.gkyl interpolate select --z0 0 plot + """ + ctx.obj = DataSpace() + + +for _command in COMMANDS: + cli.add_command(_command) + +__all__ = [ + "COMMANDS", + "COMMAND_ALIASES", + "COMMAND_SECTIONS", + "MODELS", + "PgkylGroup", + "cli", +] + +if __name__ == "__main__": + cli() diff --git a/src/postgkyl/cli/compiler.py b/src/postgkyl/cli/compiler.py new file mode 100644 index 00000000..b8a9d18f --- /dev/null +++ b/src/postgkyl/cli/compiler.py @@ -0,0 +1,759 @@ +"""Strict callable-to-command compilation and generic pipeline execution.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +import inspect +from pathlib import Path +import types +from typing import Annotated, Any, Literal, Union, get_args, get_origin, get_type_hints + +import click + +from postgkyl.cli_spec import ( + ChoiceProvider, + CliArgument, + CliType, + CommandSpec, + DatasetRef, + Execution, + KeyValue, + PipelineInput, + ResultPolicy, + Section, + command_spec, +) + +from .docstrings import parse_docstring + + +class CommandCompilationError(ValueError): + """A public callable cannot be represented by the closed CLI contract.""" + + +class CodecKind(Enum): + STRING = "string" + INTEGER = "integer" + FLOAT = "float" + BOOLEAN = "boolean" + PATH = "path" + CHOICE = "choice" + ENUM = "enum" + SEQUENCE = "sequence" + TUPLE = "tuple" + MAPPING = "mapping" + + +@dataclass(frozen=True) +class TypeCodec: + kind: CodecKind + python_type: object + optional: bool = False + choices: tuple[object, ...] = () + items: tuple["TypeCodec", ...] = () + multiple: bool = False + nargs: int = 1 + + +@dataclass(frozen=True) +class ParameterModel: + name: str + kind: inspect._ParameterKind + annotation: object + codec: TypeCodec | None + help: str | None + required: bool + default: object + injected: bool = False + dataset_ref: bool = False + argument: bool = False + + +@dataclass(frozen=True) +class CommandModel: + name: str + callable: object + canonical: object + qualname: str + spec: CommandSpec + help: str + long_help: str + parameters: tuple[ParameterModel, ...] + + @property + def section(self) -> Section: + return self.spec.section + + +class _DocumentedArgument(click.Argument): + """Click argument which retains its API parameter's help text.""" + + def __init__(self, param_decls, *, help: str | None = None, **attrs): + super().__init__(param_decls, **attrs) + self.help = help + + def get_help_record(self, ctx): + if not self.help: + return None + return self.make_metavar(ctx), self.help + + +class _GeneratedCommand(click.Command): + """Render positional parameter docs separately from option docs.""" + + def parse_args(self, ctx, args): + """Let generated boolean options omit their otherwise optional value.""" + boolean_options = { + option_name + for parameter in self.get_params(ctx) + if isinstance(parameter, _OptionalBooleanOption) + for option_name in parameter.opts + } + normalized = [] + for index, value in enumerate(args): + normalized.append(value) + option_name = value.partition("=")[0] + if value != option_name or option_name not in boolean_options: + continue + next_value = args[index + 1] if index + 1 < len(args) else None + if next_value is None or not _is_boolean_literal(next_value): + normalized.append("True") + return super().parse_args(ctx, normalized) + + def format_options(self, ctx, formatter) -> None: + arguments = [] + options = [] + for parameter in self.get_params(ctx): + record = parameter.get_help_record(ctx) + if record is None: + continue + target = arguments if isinstance(parameter, click.Argument) else options + target.append(record) + if arguments: + with formatter.section("Arguments"): + formatter.write_dl(arguments) + if options: + with formatter.section("Options"): + formatter.write_dl(options) + + +class _OptionalBooleanOption(click.Option): + """Marker for a BOOL option whose explicit value may be omitted.""" + + +_NONE_TYPE = type(None) +_SCALARS = { + str: CodecKind.STRING, + int: CodecKind.INTEGER, + float: CodecKind.FLOAT, + bool: CodecKind.BOOLEAN, + Path: CodecKind.PATH, +} +_RESERVED_SHORT_OPTIONS = frozenset({"h"}) +_BOOLEAN_LITERALS = frozenset( + {"1", "0", "yes", "no", "true", "false", "on", "off", "t", "f", "y", "n"}) + + +def _is_boolean_literal(value: str) -> bool: + return value.strip().lower() in _BOOLEAN_LITERALS + + +def _short_option_names( + parameters: tuple[ParameterModel, ...]) -> dict[str, str]: + """Assign each available initial to its first exposed parameter.""" + assigned: dict[str, str] = {} + claimed = set(_RESERVED_SHORT_OPTIONS) + for parameter in parameters: + if parameter.injected or parameter.argument: + continue + initial = parameter.name[0] + if initial in claimed: + continue + assigned[parameter.name] = f"-{initial}" + claimed.add(initial) + return assigned + + +def _error(fn, parameter: str, message: str) -> CommandCompilationError: + qualname = f"{getattr(fn, '__module__', '')}.{getattr(fn, '__qualname__', fn)}" + return CommandCompilationError( + f"{qualname}: parameter {parameter!r}: {message}") + + +def _unwrap_annotated(annotation): + markers: list[object] = [] + while get_origin(annotation) is Annotated: + args = get_args(annotation) + annotation = args[0] + markers.extend(args[1:]) + return annotation, tuple(markers) + + +def _optional(annotation): + origin = get_origin(annotation) + if origin not in (Union, types.UnionType): + return annotation, False + args = get_args(annotation) + non_none = tuple(arg for arg in args if arg is not _NONE_TYPE) + if len(non_none) == 1 and len(non_none) != len(args): + return non_none[0], True + return annotation, False + + +def _codec(fn, name: str, annotation, markers: tuple[object, ...]) -> TypeCodec: + providers = [ + marker for marker in markers if isinstance(marker, ChoiceProvider) + ] + key_values = [marker for marker in markers if isinstance(marker, KeyValue)] + cli_types = [marker for marker in markers if isinstance(marker, CliType)] + unknown = [ + marker for marker in markers + if not isinstance(marker, (ChoiceProvider, CliArgument, CliType, KeyValue, + DatasetRef, PipelineInput)) + ] + if unknown: + raise _error(fn, name, f"unsupported Annotated marker {unknown[0]!r}") + if len(providers) > 1 or len(key_values) > 1 or len(cli_types) > 1: + raise _error(fn, name, "duplicate Annotated codec marker") + if providers and key_values: + raise _error(fn, name, "ChoiceProvider and KeyValue cannot be combined") + if cli_types: + annotation = cli_types[0].annotation + annotation, optional = _optional(annotation) + if annotation in (Any, inspect.Parameter.empty): + raise _error(fn, name, "missing or Any annotation") + + if providers: + try: + provided = providers[0].provider() + if isinstance(provided, (str, bytes)): + raise TypeError("provider returned text instead of a choice collection") + values = tuple(provided) + except Exception as exc: + raise _error(fn, name, f"choice provider failed: {exc}") from exc + if not values: + raise _error(fn, name, "choice provider returned no choices") + if annotation not in _SCALARS: + raise _error(fn, name, "choice provider requires a scalar annotation") + if not all( + isinstance(value, annotation) + and not (annotation is int and isinstance(value, bool)) + for value in values): + raise _error( + fn, name, + "choice provider values do not match the annotated scalar type") + if len(set(values)) != len(values): + raise _error(fn, name, "choice provider returned duplicate choices") + return TypeCodec(CodecKind.CHOICE, annotation, optional, choices=values) + if key_values and get_origin(annotation) not in (dict, Mapping): + raise _error(fn, name, "KeyValue requires a mapping annotation") + if annotation in _SCALARS: + return TypeCodec(_SCALARS[annotation], annotation, optional) + if inspect.isclass(annotation) and issubclass(annotation, Enum): + values = tuple(member.value for member in annotation) + if not all(isinstance(value, (str, int, float)) for value in values): + raise _error(fn, name, "Enum values must be CLI scalars") + return TypeCodec(CodecKind.ENUM, annotation, optional, choices=values) + origin = get_origin(annotation) + args = get_args(annotation) + if origin is Literal: + if not args or not all( + isinstance(value, (str, int, float)) for value in args): + raise _error(fn, name, "Literal choices must be strings or numbers") + return TypeCodec(CodecKind.CHOICE, + annotation, + optional, + choices=tuple(args)) + if origin is list: + if len(args) != 1: + raise _error(fn, name, "list must have exactly one item type") + item = _codec(fn, name, args[0], ()) + return TypeCodec(CodecKind.SEQUENCE, + annotation, + optional, + items=(item, ), + multiple=True) + if origin is tuple: + if not args: + raise _error(fn, name, "tuple must declare its item type(s)") + if len(args) == 2 and args[1] is Ellipsis: + item = _codec(fn, name, args[0], ()) + return TypeCodec(CodecKind.SEQUENCE, + annotation, + optional, + items=(item, ), + multiple=True) + items = tuple(_codec(fn, name, arg, ()) for arg in args) + return TypeCodec(CodecKind.TUPLE, + annotation, + optional, + items=items, + nargs=len(items)) + if origin in (dict, Mapping): + if not key_values: + raise _error(fn, name, "mapping needs Annotated[..., KeyValue()]") + if len(args) != 2: + raise _error(fn, name, "mapping must declare key and value types") + key = _codec(fn, name, args[0], ()) + value = _codec(fn, name, args[1], ()) + composite = {CodecKind.SEQUENCE, CodecKind.TUPLE, CodecKind.MAPPING} + if key.kind in composite or value.kind in composite: + raise _error(fn, name, "mapping keys and values must be CLI scalars") + return TypeCodec(CodecKind.MAPPING, + annotation, + optional, + items=(key, value), + multiple=True) + if get_origin(annotation) in (Union, types.UnionType): + raise _error( + fn, name, + f"unsupported union {annotation!r}; only T | None is lossless") + raise _error(fn, name, f"unsupported annotation {annotation!r}") + + +def _type_hints(fn) -> dict[str, object]: + annotations = dict(getattr(fn, "__annotations__", {})) + deferred = { + name: annotation + for name, annotation in annotations.items() + if isinstance(annotation, str) + } + if not deferred: + return annotations + try: + source = types.SimpleNamespace(__annotations__=deferred) + resolved = get_type_hints(source, + globalns=getattr(fn, "__globals__", None), + include_extras=True) + except Exception as exc: + qualname = f"{getattr(fn, '__module__', '')}.{getattr(fn, '__qualname__', fn)}" + raise CommandCompilationError( + f"{qualname}: annotations could not be resolved: {exc}") from exc + # Concrete annotations installed by an API owner are already resolved. + # Evaluating them again is observably broken for + # ``Annotated[T | None, ...]`` on Python 3.10, so only deferred strings pass + # through ``get_type_hints`` and concrete objects remain authoritative. + for name, annotation in annotations.items(): + if isinstance(annotation, str): + annotations[name] = resolved[name] + return annotations + + +def compile_callable(fn, *, name: str | None = None) -> CommandModel: + """Inspect one marked callable into an immutable, Click-free model.""" + spec = command_spec(fn) + if spec is None: + raise CommandCompilationError(f"{fn!r} has no CommandSpec") + canonical = fn + signature = inspect.signature(canonical) + hints = _type_hints(canonical) + raw: list[tuple[inspect.Parameter, object, tuple[object, ...], bool, bool, + bool]] = [] + for index, parameter in enumerate(signature.parameters.values()): + if parameter.kind is inspect.Parameter.VAR_KEYWORD: + raise _error(canonical, parameter.name, "**kwargs is not representable") + annotation = hints.get(parameter.name, parameter.annotation) + base, markers = _unwrap_annotated(annotation) + marker_injected = any( + isinstance(marker, PipelineInput) for marker in markers) + if sum(isinstance(marker, PipelineInput) for marker in markers) > 1: + raise _error(canonical, parameter.name, "duplicate PipelineInput marker") + is_receiver = parameter.name == "self" or ( + index == 0 and spec.execution + in (Execution.MAP_REPLACE, Execution.MAP_APPEND, + Execution.MAP_OR_TERMINAL_EACH, Execution.TERMINAL_EACH)) + is_variadic_input = parameter.kind is inspect.Parameter.VAR_POSITIONAL and ( + spec.execution in (Execution.COMBINE, Execution.TERMINAL_ALL)) + injected = marker_injected or is_receiver or is_variadic_input + if injected and parameter.name != "self" and base in ( + Any, inspect.Parameter.empty): + raise _error(canonical, parameter.name, + "pipeline receivers need a concrete annotation") + if parameter.kind is inspect.Parameter.VAR_POSITIONAL and not injected: + raise _error(canonical, parameter.name, + "*args is not a declared pipeline receiver") + is_dataset_ref = any(isinstance(marker, DatasetRef) for marker in markers) + if sum(isinstance(marker, DatasetRef) for marker in markers) > 1: + raise _error(canonical, parameter.name, "duplicate DatasetRef marker") + if is_dataset_ref and injected: + raise _error(canonical, parameter.name, + "cannot be both DatasetRef and PipelineInput") + is_argument = any(isinstance(marker, CliArgument) for marker in markers) + if sum(isinstance(marker, CliArgument) for marker in markers) > 1: + raise _error(canonical, parameter.name, "duplicate CliArgument marker") + if is_argument and injected: + raise _error(canonical, parameter.name, + "cannot be both CliArgument and PipelineInput") + if is_argument and is_dataset_ref: + raise _error(canonical, parameter.name, + "cannot be both CliArgument and DatasetRef") + if is_argument and parameter.kind not in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD): + raise _error(canonical, parameter.name, + "CliArgument requires a positional Python parameter") + raw.append( + (parameter, base, markers, injected, is_dataset_ref, is_argument)) + + injected_parameters = [item for item in raw if item[3]] + if spec.execution is Execution.LOAD and injected_parameters: + raise CommandCompilationError( + f"{canonical.__module__}.{canonical.__qualname__}: " + "LOAD commands cannot declare a pipeline receiver") + docs = parse_docstring(canonical, + required=set(signature.parameters), + signature_names=set(signature.parameters)) + models: list[ParameterModel] = [] + for parameter, base, markers, injected, is_dataset_ref, is_argument in raw: + required = parameter.default is inspect.Parameter.empty + default = None if required else parameter.default + codec = None + if not injected: + codec = (TypeCodec(CodecKind.STRING, str, + optional=not required) if is_dataset_ref else _codec( + canonical, parameter.name, base, markers)) + if codec.kind is CodecKind.BOOLEAN and parameter.default is not False: + raise _error(canonical, parameter.name, + "boolean CLI options must default to False") + if is_argument and (codec.multiple or codec.nargs != 1): + raise _error(canonical, parameter.name, + "CliArgument currently supports scalar values only") + models.append( + ParameterModel(name=parameter.name, + kind=parameter.kind, + annotation=base, + codec=codec, + help=docs.parameters.get(parameter.name), + required=required, + default=default, + injected=injected, + dataset_ref=is_dataset_ref, + argument=is_argument)) + + command_name = name or getattr(canonical, "__name__", "") + if not command_name or command_name.startswith("-"): + raise CommandCompilationError( + f"{canonical!r}: invalid command name {command_name!r}") + qualname = f"{canonical.__module__}.{canonical.__qualname__}" + return CommandModel(command_name, fn, canonical, qualname, spec, docs.summary, + docs.long_help, tuple(models)) + + +def compile_public_surface(callables) -> tuple[CommandModel, ...]: + """Compile and validate a complete discovered surface atomically.""" + models = tuple( + compile_callable(item.callable, name=item.name) for item in callables) + seen: dict[str, CommandModel] = {} + for model in models: + previous = seen.get(model.name) + if previous is not None and previous.canonical is not model.canonical: + raise CommandCompilationError( + f"command name collision {model.name!r}: {previous.qualname} and {model.qualname}" + ) + seen[model.name] = model + return tuple( + sorted( + seen.values(), + key=lambda model: + (list(Section).index(model.section), model.spec.order, model.name))) + + +def _click_scalar(codec: TypeCodec): + if codec.kind is CodecKind.STRING: + return click.STRING + if codec.kind is CodecKind.INTEGER: + return click.INT + if codec.kind is CodecKind.FLOAT: + return click.FLOAT + if codec.kind is CodecKind.BOOLEAN: + return click.BOOL + if codec.kind is CodecKind.PATH: + return click.Path(path_type=Path) + if codec.kind in (CodecKind.CHOICE, CodecKind.ENUM): + return click.Choice(codec.choices, case_sensitive=True) + if codec.kind in (CodecKind.SEQUENCE, CodecKind.MAPPING): + return _click_scalar( + codec.items[0]) if codec.kind is CodecKind.SEQUENCE else click.STRING + if codec.kind is CodecKind.TUPLE: + return click.Tuple([_click_scalar(item) for item in codec.items]) + raise AssertionError(codec.kind) + + +def _convert_scalar(value, codec: TypeCodec): + if value is None: + return None + if codec.kind is CodecKind.ENUM: + return codec.python_type(value) + return value + + +def _convert(value, codec: TypeCodec): + if value is None: + return None + if codec.kind is CodecKind.SEQUENCE: + return [_convert_scalar(item, codec.items[0]) for item in value] + if codec.kind is CodecKind.TUPLE: + return tuple( + _convert_scalar(item, sub) for item, sub in zip(value, codec.items)) + if codec.kind is CodecKind.MAPPING: + result = {} + key_codec, value_codec = codec.items + for entry in value: + key, separator, item = entry.partition("=") + if not separator or not key: + raise click.BadParameter("expected key=value") + click_key = _click_scalar(key_codec).convert(key, None, None) + click_value = _click_scalar(value_codec).convert(item, None, None) + if click_key in result: + raise click.BadParameter(f"duplicate mapping key {click_key!r}") + result[_convert_scalar(click_key, key_codec)] = _convert_scalar( + click_value, value_codec) + return result or None if codec.optional else result + return _convert_scalar(value, codec) + + +def _resolve_tag(ctx, parameter: str, tag: str): + matches = [dataset for dataset in ctx.obj.datasets if dataset.tag == tag] + if not matches: + raise click.UsageError(f"--{parameter}: no dataset tagged {tag!r}") + if len(matches) != 1: + raise click.UsageError( + f"--{parameter}: tag {tag!r} matches {len(matches)} datasets") + return matches[0] + + +def _is_dataset(value) -> bool: + return all(hasattr(value, name) for name in ("ctx", "grid", "values")) + + +def _datasets_from_result(result) -> list: + members = getattr(result, "datasets", None) + if members is not None: + return list(members) + if _is_dataset(result): + return [result] + if isinstance(result, + (list, tuple)) and all(_is_dataset(item) for item in result): + return list(result) + return [] + + +def _present(value) -> None: + if value is None: + return + if isinstance(value, (list, tuple)): + for item in value: + _present(item) + return + if not _is_dataset(value): + click.echo(value) + + +def _selected(ctx) -> list: + return list(ctx.obj.datasets) + + +def _call(model: CommandModel, selected: list, values: dict, ctx): + args: list[object] = [] + kwargs: dict[str, object] = {} + referenced: list[object] = [] + for parameter in model.parameters: + if parameter.injected: + if parameter.kind is inspect.Parameter.VAR_POSITIONAL: + args.extend(selected) + elif parameter.name == "self": + if len(selected) != 1: + raise click.UsageError( + f"{model.name}: expected exactly one pipeline input") + args.append(selected[0]) + elif model.spec.execution in (Execution.MAP_REPLACE, Execution.MAP_APPEND, + Execution.MAP_OR_TERMINAL_EACH, + Execution.TERMINAL_EACH): + args.append(selected[0]) + elif parameter.kind in (inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD): + args.append(selected) + else: + kwargs[parameter.name] = selected + continue + value = values[parameter.name] + if parameter.dataset_ref: + if value is not None: + value = _resolve_tag(ctx, parameter.name, value) + referenced.append(value) + elif parameter.codec is not None: + value = _convert(value, parameter.codec) + if parameter.kind in (inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD): + args.append(value) + else: + kwargs[parameter.name] = value + return model.canonical(*args, **kwargs), referenced + + +def execute_model(ctx, model: CommandModel, values: dict): + """Invoke one model through its generic working-set execution adapter.""" + selected = _selected(ctx) + execution = model.spec.execution + needs_input = execution is not Execution.LOAD + if needs_input and not selected and not any(p.dataset_ref + for p in model.parameters): + raise click.UsageError(f"{model.name}: no datasets selected") + + try: + if execution in (Execution.MAP_REPLACE, Execution.MAP_APPEND, + Execution.MAP_OR_TERMINAL_EACH, Execution.TERMINAL_EACH): + outputs: list[object] = [] + for dataset in selected: + result, _ = _call(model, [dataset], values, ctx) + outputs.append(result) + if execution is Execution.MAP_REPLACE: + replacements = iter(outputs) + chosen = {id(dataset) for dataset in selected} + ctx.obj.datasets = [ + next(replacements) if id(dataset) in chosen else dataset + for dataset in ctx.obj.datasets + ] + elif execution is Execution.MAP_APPEND: + if model.spec.consumes_inputs: + consumed = {id(dataset) for dataset in selected} + ctx.obj.datasets = [ + dataset for dataset in ctx.obj.datasets + if id(dataset) not in consumed + ] + for result in outputs: + ctx.obj.datasets.extend(_datasets_from_result(result)) + elif execution is Execution.MAP_OR_TERMINAL_EACH: + by_input = { + id(dataset): result + for dataset, result in zip(selected, outputs) + } + ctx.obj.datasets = [ + by_input[id(dataset)] if id(dataset) in by_input + and _is_dataset(by_input[id(dataset)]) else dataset + for dataset in ctx.obj.datasets + ] + if model.spec.result is ResultPolicy.VALUE: + for result in outputs: + _present(result) + else: + if model.spec.result is ResultPolicy.VALUE: + for result in outputs: + _present(result) + return outputs + + result, referenced = _call(model, selected, values, ctx) + if execution is Execution.LOAD: + ctx.obj.datasets.extend(_datasets_from_result(result)) + if model.spec.result is ResultPolicy.VALUE: + _present(result) + elif execution is Execution.COMBINE: + result_datasets = _datasets_from_result(result) + if len(result_datasets) == len(selected) and sorted(map(id, result_datasets)) \ + == sorted(map(id, selected)): + ordered = iter(result_datasets) + selected_ids = {id(dataset) for dataset in selected} + ctx.obj.datasets = [ + next(ordered) if id(dataset) in selected_ids else dataset + for dataset in ctx.obj.datasets + ] + return result + if model.spec.consumes_inputs: + consumed = {id(dataset) for dataset in (referenced or selected)} + ctx.obj.datasets = [ + dataset for dataset in ctx.obj.datasets + if id(dataset) not in consumed + ] + ctx.obj.datasets.extend(result_datasets) + if model.spec.result is ResultPolicy.VALUE: + _present(result) + elif execution is Execution.TERMINAL_ALL: + if model.spec.result is ResultPolicy.VALUE: + _present(result) + return result + except click.ClickException: + raise + except (ValueError, TypeError, OSError) as exc: + raise click.UsageError(str(exc)) from exc + + +def build_click_command(model: CommandModel, + *, + command_class=_GeneratedCommand) -> click.Command: + """Lower an immutable model to a Click command.""" + params: list[click.Parameter] = [] + short_options = _short_option_names(model.parameters) + for parameter in model.parameters: + if parameter.injected: + continue + codec = parameter.codec + assert codec is not None + if parameter.argument: + attrs = dict(type=_click_scalar(codec), required=parameter.required) + if not parameter.required: + default = parameter.default + if codec.kind is CodecKind.ENUM and isinstance(default, Enum): + default = default.value + attrs["default"] = default + params.append( + _DocumentedArgument([parameter.name], help=parameter.help, **attrs)) + continue + attrs = dict( + type=_click_scalar(codec), + required=parameter.required, + help=parameter.help, + show_default=not parameter.required, + multiple=codec.multiple, + ) + if not parameter.required: + default = parameter.default + if codec.kind is CodecKind.ENUM and isinstance(default, Enum): + default = default.value + if codec.multiple and default is None: + default = () + attrs["default"] = default + declarations = [f"--{parameter.name}"] + if parameter.name in short_options: + declarations.append(short_options[parameter.name]) + declarations.append(parameter.name) + option_class = (_OptionalBooleanOption + if codec.kind is CodecKind.BOOLEAN else click.Option) + if codec.kind is CodecKind.BOOLEAN: + attrs["metavar"] = "[BOOLEAN]" + params.append(option_class(declarations, **attrs)) + + @click.pass_context + def callback(click_context, **kwargs): + return execute_model(click_context, model, kwargs) + + return command_class(model.name, + params=params, + callback=callback, + help=model.long_help, + short_help=model.help) + + +def group_by_section(models: tuple[CommandModel, ...]) -> dict[str, list[str]]: + """Derive the flat help presentation from compiled models.""" + return { + section.value: + [model.name for model in models if model.section is section] + for section in Section + if any(model.section is section for model in models) + } + + +__all__ = [ + "CodecKind", + "CommandCompilationError", + "CommandModel", + "ParameterModel", + "TypeCodec", + "build_click_command", + "compile_callable", + "compile_public_surface", + "execute_model", + "group_by_section", +] diff --git a/src/postgkyl/cli/discovery.py b/src/postgkyl/cli/discovery.py new file mode 100644 index 00000000..e8c1670c --- /dev/null +++ b/src/postgkyl/cli/discovery.py @@ -0,0 +1,135 @@ +"""Deterministic discovery of command metadata from public API roots.""" + +from __future__ import annotations + +from dataclasses import dataclass +import inspect +from types import ModuleType + +import postgkyl +from postgkyl.cli_spec import command_spec, hidden_spec + + +class SurfaceClassificationError(ValueError): + """A public CLI candidate is neither exposed nor explicitly hidden.""" + + +@dataclass(frozen=True) +class DiscoveredCallable: + name: str + callable: object + public_path: str + + +def _classify(obj, path: str) -> bool: + exposed = command_spec(obj) is not None + excluded = hidden_spec(obj) is not None + if exposed == excluded: + state = "both exposed and hidden" if exposed else "unclassified" + raise SurfaceClassificationError(f"{path}: public callable is {state}") + return exposed + + +def _functions(module: ModuleType): + public = getattr(module, "__all__", None) + names = public if public is not None else sorted( + name for name, value in vars(module).items() if not name.startswith("_") + and inspect.isfunction(value) and value.__module__ == module.__name__) + for name in names: + value = getattr(module, name) + if inspect.isfunction(value): + yield name, value + + +def _diagnostic_modules(root: ModuleType): + seen: set[int] = set() + + def visit(module): + if id(module) in seen: + return + seen.add(id(module)) + yield module + for name in getattr(module, "__all__", ()): + value = getattr(module, name) + if isinstance( + value, + ModuleType) and value.__name__.startswith("postgkyl.diagnostics"): + yield from visit(value) + + yield from visit(root) + + +def discover_public_surface(facade=postgkyl) -> tuple[DiscoveredCallable, ...]: + """Walk the declared public roots and return all exposed callables.""" + found: list[DiscoveredCallable] = [] + classified: set[tuple[int, str]] = set() + + def consider(obj, path: str, name: str) -> None: + key = (id(obj), name) + if key in classified: + return + classified.add(key) + if _classify(obj, path): + found.append(DiscoveredCallable(name, obj, path)) + + # Fluent methods are the authoritative inventory of per-dataset commands. + for cls_name in ("GData", "GDataGroup"): + cls = getattr(facade, cls_name) + for name, value in sorted(cls.__dict__.items()): + if name.startswith("_") or not callable(value): + continue + consider(value, f"{cls.__module__}.{cls.__qualname__}.{name}", name) + + # Public facade functions not already reached through their fluent view. + for name in facade.__all__: + value = getattr(facade, name) + if inspect.isfunction(value): + path = f"postgkyl.{name}" + if value.__module__.startswith("postgkyl.diagnostics"): + _classify(value, path) + else: + consider(value, path, name) + + diagnostics = facade.diagnostics + for module in _diagnostic_modules(diagnostics): + if module is diagnostics: + continue + relative = module.__name__.removeprefix("postgkyl.diagnostics.") + # Model-family packages group related diagnostic modules without erasing + # each module's public command vocabulary. For example, + # diagnostics.mom.five_moment.pressure becomes ``five_moment_pressure`` + # and cannot collide with ten_moment.pressure. + namespace = relative.rsplit(".", 1)[-1] + for name, value in _functions(module): + if not value.__module__.startswith("postgkyl.diagnostics"): + # Compatibility re-exports owned by a lower layer keep that owner's + # canonical command (for example operations.gyrokinetics.gk_rz -> + # ``gk_rz``); the diagnostic alias is classified but does not invent + # a second command or move it into the wrong help section. + _classify(value, f"{module.__name__}.{name}") + continue + consider(value, f"{module.__name__}.{name}", f"{namespace}_{name}") + + # Registry values are public vocabulary roots too. Identity de-duplication + # means aliases do not invent additional command names. + variables = getattr(module, "VARIABLES", None) + if isinstance(variables, dict): + for value in variables.values(): + if inspect.isfunction(value): + consider(value, f"{module.__name__}.VARIABLES[{value.__name__!r}]", + f"{namespace}_{value.__name__}") + + # A fluent view and facade function may be aliases of the same operation. + # Keep one deterministic view. + unique: dict[tuple[str, int], DiscoveredCallable] = {} + for item in found: + key = (item.name, id(item.callable)) + unique.setdefault(key, item) + return tuple( + sorted(unique.values(), key=lambda item: (item.name, item.public_path))) + + +__all__ = [ + "DiscoveredCallable", "SurfaceClassificationError", + "discover_public_surface" +] diff --git a/src/postgkyl/cli/docstrings.py b/src/postgkyl/cli/docstrings.py new file mode 100644 index 00000000..8a1ae788 --- /dev/null +++ b/src/postgkyl/cli/docstrings.py @@ -0,0 +1,109 @@ +"""Strict parser for the Google-style subset used by generated commands.""" + +from __future__ import annotations + +from dataclasses import dataclass +import inspect +import re + + +class DocstringError(ValueError): + """A canonical callable's documentation cannot be lowered losslessly.""" + + +@dataclass(frozen=True) +class ParsedDocstring: + summary: str + long_help: str + parameters: dict[str, str] + + +_SECTION = re.compile(r"^([A-Za-z][A-Za-z ]*):\s*$") +_ARG = re.compile(r"^\s{2,}(\*{0,2}[A-Za-z_]\w*)(?:\s*\([^)]*\))?:\s*(.*)$") + + +def _paragraph(lines: list[str]) -> str: + out: list[str] = [] + for line in lines: + if not line.strip() or _SECTION.match(line): + if out: + break + continue + out.append(line.strip()) + return " ".join(out) + + +def _narrative(lines: list[str]) -> str: + """Return prose preceding the first structured docstring section.""" + end = next( + (index for index, line in enumerate(lines) if _SECTION.match(line)), + len(lines)) + return "\n".join(lines[:end]).strip() + + +def parse_docstring(obj, + *, + required: set[str] | None = None, + signature_names: set[str] | None = None) -> ParsedDocstring: + """Parse and validate one canonical callable's command documentation.""" + qualname = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', obj)}" + doc = inspect.getdoc(obj) + if not doc: + raise DocstringError(f"{qualname}: missing docstring") + lines = doc.splitlines() + summary = _paragraph(lines) + if not summary: + raise DocstringError(f"{qualname}: missing first-paragraph description") + + entries: dict[str, str] = {} + args_sections = [i for i, line in enumerate(lines) if line == "Args:"] + if len(args_sections) > 1: + raise DocstringError(f"{qualname}: duplicate Args sections") + args_at = args_sections[0] if args_sections else None + if args_at is not None: + current: str | None = None + chunks: list[str] = [] + + def finish() -> None: + if current is None: + return + text = " ".join(part for part in chunks if part).strip() + if not text: + raise DocstringError( + f"{qualname}: parameter {current!r} has no description") + if current in entries: + raise DocstringError( + f"{qualname}: parameter {current!r} is documented twice") + entries[current] = text + + for line in lines[args_at + 1:]: + if line and not line[0].isspace() and _SECTION.match(line): + break + match = _ARG.match(line) + if match: + finish() + current = match.group(1).lstrip("*") + chunks = [match.group(2).strip()] + elif current is not None and (not line.strip() + or line.startswith(" ")): + chunks.append(line.strip()) + elif line.strip(): + raise DocstringError( + f"{qualname}: malformed Args entry: {line.strip()!r}") + finish() + + signature_names = signature_names or set() + unknown = set(entries) - signature_names + if unknown: + names = ", ".join(sorted(unknown)) + raise DocstringError( + f"{qualname}: documented parameter(s) absent from signature: {names}") + for name in sorted(required or ()): + if name not in entries: + raise DocstringError(f"{qualname}: parameter {name!r} is undocumented") + return ParsedDocstring(summary=summary, + long_help=_narrative(lines), + parameters=entries) + + +__all__ = ["DocstringError", "ParsedDocstring", "parse_docstring"] diff --git a/src/postgkyl/cli/state.py b/src/postgkyl/cli/state.py new file mode 100644 index 00000000..c1a53da5 --- /dev/null +++ b/src/postgkyl/cli/state.py @@ -0,0 +1,18 @@ +"""Shared CLI state -- the chained pipeline's scratch space (``ctx.obj``).""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class DataSpace: + """Datasets flowing through a chained command line. + + ``datasets`` is the working set every generated verb transforms. + """ + + datasets: list = field(default_factory=list) + + def __iter__(self): + return iter(self.datasets) diff --git a/src/postgkyl/cli_spec.py b/src/postgkyl/cli_spec.py new file mode 100644 index 00000000..86df52e7 --- /dev/null +++ b/src/postgkyl/cli_spec.py @@ -0,0 +1,192 @@ +"""Frozen, dependency-free records describing the generated CLI surface. + +This package deliberately knows nothing about Click or postgkyl datasets. +API-owning modules attach these records to their public callables; the CLI is +the only layer that interprets them. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum, auto +from typing import Callable + + +class Section(Enum): + """Top-level help sections, in display order.""" + + VERBS = "Verbs" + DIAGNOSTICS = "Diagnostics" + RENDER = "Render" + UTILITY = "Utility" + + +class Execution(Enum): + """Closed set of working-set adapters used by the CLI runtime.""" + + MAP_REPLACE = auto() + MAP_APPEND = auto() + MAP_OR_TERMINAL_EACH = auto() + COMBINE = auto() + TERMINAL_EACH = auto() + TERMINAL_ALL = auto() + LOAD = auto() + + +class ResultPolicy(Enum): + """How a returned API value is presented at the command-line edge.""" + + DATA = auto() + VALUE = auto() + SILENT = auto() + + +@dataclass(frozen=True) +class CommandSpec: + """Pipeline metadata attached to one canonical API callable.""" + + section: Section + execution: Execution + consumes_inputs: bool = False + result: ResultPolicy = ResultPolicy.DATA + order: int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.section, Section) or not isinstance( + self.execution, Execution) or not isinstance(self.result, ResultPolicy): + raise TypeError( + "CommandSpec section, execution, and result must be enums") + if not isinstance(self.consumes_inputs, bool): + raise TypeError("CommandSpec consumes_inputs must be a bool value") + if not isinstance(self.order, int) or isinstance(self.order, bool): + raise TypeError("CommandSpec order must be an integer") + if self.execution is Execution.LOAD: + if self.consumes_inputs: + raise ValueError("LOAD commands cannot consume working-set inputs") + if self.consumes_inputs and self.execution not in (Execution.MAP_APPEND, + Execution.COMBINE): + raise ValueError( + "only MAP_APPEND and COMBINE commands may consume their inputs") + if self.execution in (Execution.TERMINAL_EACH, Execution.TERMINAL_ALL): + if self.consumes_inputs: + raise ValueError("terminal commands cannot consume working-set inputs") + if self.result is ResultPolicy.DATA: + raise ValueError("terminal commands need a non-DATA result policy") + + +@dataclass(frozen=True) +class DatasetRef: + """Resolve this API dataset parameter from a unique working-set tag.""" + + default_tag: str | None = None + + def __post_init__(self) -> None: + if self.default_tag is not None and not self.default_tag: + raise ValueError("DatasetRef default_tag must be non-empty when supplied") + + +@dataclass(frozen=True) +class PipelineInput: + """Inject this parameter from the selected CLI working set.""" + + +@dataclass(frozen=True) +class CliType: + """Lossless CLI input type for a broader direct-Python annotation.""" + + annotation: object + + +@dataclass(frozen=True) +class CliArgument: + """Expose this API parameter as a positional command-line argument.""" + + +@dataclass(frozen=True) +class ChoiceProvider: + """Obtain an option's choices from the API registry returned by provider.""" + + provider: Callable[[], object] + + def __post_init__(self) -> None: + if not callable(self.provider): + raise TypeError("ChoiceProvider requires a callable provider") + + +@dataclass(frozen=True) +class KeyValue: + """Parse a mapping as repeated ``key=value`` option values.""" + + +@dataclass(frozen=True) +class CliHidden: + """Explicitly exclude a public callable from command generation.""" + + reason: str + + def __post_init__(self) -> None: + if not self.reason.strip(): + raise ValueError("CliHidden requires a non-empty reason") + + +_COMMAND_ATTR = "__postgkyl_command_spec__" +_HIDDEN_ATTR = "__postgkyl_cli_hidden__" + + +def command(spec: CommandSpec): + """Attach ``spec`` to a callable without wrapping or registering it.""" + if not isinstance(spec, CommandSpec): + raise TypeError("command() requires a CommandSpec") + + def decorate(fn): + if hasattr(fn, _HIDDEN_ATTR): + raise ValueError(f"{fn!r} is already marked CliHidden") + previous = getattr(fn, _COMMAND_ATTR, None) + if previous is not None and previous != spec: + raise ValueError(f"{fn!r} already has a different CommandSpec") + setattr(fn, _COMMAND_ATTR, spec) + return fn + + return decorate + + +def hidden(reason: str): + """Explicitly mark a public callable as unavailable from the CLI.""" + marker = CliHidden(reason) + + def decorate(fn): + if hasattr(fn, _COMMAND_ATTR): + raise ValueError(f"{fn!r} already has a CommandSpec") + setattr(fn, _HIDDEN_ATTR, marker) + return fn + + return decorate + + +def command_spec(fn) -> CommandSpec | None: + """Return the immutable command metadata attached to ``fn``.""" + return getattr(fn, _COMMAND_ATTR, None) + + +def hidden_spec(fn) -> CliHidden | None: + """Return the explicit CLI exclusion attached to ``fn``.""" + return getattr(fn, _HIDDEN_ATTR, None) + + +__all__ = [ + "ChoiceProvider", + "CliArgument", + "CliHidden", + "CliType", + "CommandSpec", + "DatasetRef", + "Execution", + "KeyValue", + "PipelineInput", + "ResultPolicy", + "Section", + "command", + "command_spec", + "hidden", + "hidden_spec", +] diff --git a/src/postgkyl/commands/__init__.py b/src/postgkyl/commands/__init__.py deleted file mode 100644 index 6339c117..00000000 --- a/src/postgkyl/commands/__init__.py +++ /dev/null @@ -1,58 +0,0 @@ -from postgkyl.commands.data_space import DataSpace -from postgkyl.commands.config import config - -from postgkyl.commands import ev_cmd - -from postgkyl.commands.agyro import agyro -from postgkyl.commands.agyro import mom_agyro -from postgkyl.commands.animate import animate -from postgkyl.commands.bparrotate import bparrotate -from postgkyl.commands.bperprotate import bperprotate -from postgkyl.commands.collect import collect -from postgkyl.commands.current import current -from postgkyl.commands.dg_evproj import dg_evproj -from postgkyl.commands.dg_avg import dg_avg -from postgkyl.commands.differentiate import differentiate -from postgkyl.commands.energetics import energetics -from postgkyl.commands.euler import euler -from postgkyl.commands.ev import ev -from postgkyl.commands.extractinput import extractinput -from postgkyl.commands.fft import fft -from postgkyl.commands.gkyl_pkpm import pkpm -from postgkyl.commands.gk_nodes import gk_nodes -from postgkyl.commands.grid import grid -from postgkyl.commands.growth import growth -from postgkyl.commands.info import info -from postgkyl.commands.integrate import integrate -from postgkyl.commands.interpolate import interpolate -from postgkyl.commands.laguerre_compose import laguerrecompose -from postgkyl.commands.listoutputs import listoutputs -from postgkyl.commands.load import load -from postgkyl.commands.magsq import magsq -from postgkyl.commands.mask import mask -from postgkyl.commands.mhd import mhd -from postgkyl.commands.parrotate import parrotate -from postgkyl.commands.gk_energy_balance import gk_energy_balance -from postgkyl.commands.gk_distf import load_gk_distf -from postgkyl.commands.gk_distf import gk_distf -from postgkyl.commands.dg_local_poly import dg_local_poly -from postgkyl.commands.gk_load_quantity import gk_load_quantity -from postgkyl.commands.gk_particle_balance import gk_particle_balance -from postgkyl.commands.gk_rz import gk_rz -from postgkyl.commands.gk_fluxsurf import gk_fluxsurf -from postgkyl.commands.perprotate import perprotate -from postgkyl.commands.plot import plot -from postgkyl.commands.pr import pr -from postgkyl.commands.relchange import relchange -from postgkyl.commands.select import select -from postgkyl.commands.status import activate -from postgkyl.commands.status import deactivate -from postgkyl.commands.style import style -from postgkyl.commands.tenmoment import tenmoment -from postgkyl.commands.trajectory import trajectory -from postgkyl.commands.transform_frame import transformframe -from postgkyl.commands.val2coord import val2coord -from postgkyl.commands.velocity import velocity -from postgkyl.commands.write import write - -from postgkyl.commands import temp diff --git a/src/postgkyl/commands/agyro.py b/src/postgkyl/commands/agyro.py deleted file mode 100644 index a884c651..00000000 --- a/src/postgkyl/commands/agyro.py +++ /dev/null @@ -1,70 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.tools import get_agyro, get_gkyl_10m_agyro -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--measure", "-m", default="frobenius", show_default=True, - type=click.Choice(["swisdak", "frobenius"]), - help="Specify how to calculate agyrotropy.") -@click.option("--pressure", "-p", default="pressure", show_default=True, - help="Tag for input pressure.") -@click.option("--bfield", "-b", default="field", show_default=True, - help="Tag for input EM field.") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.pass_context -def agyro(ctx, **kwargs): - """Compute a measure of agyrotropy. - - Default measure is taken from Swisdak 2015. Optionally computes agyrotropy as - Frobenius norm of agyrotropic pressure tensor. - """ - verb_print(ctx, "Starting agyro") - - data = ctx.obj["data"] - tag = "agyro" - if kwargs["tag"]: - tag = kwargs["tag"] - # end - - for pressure, bfield in zip(data.iterator(kwargs["pressure"]), data.iterator(kwargs["bfield"])): - grid, agyro_vals = get_agyro(p_in=pressure, b_in=bfield, measure=kwargs["measure"]) - out = GData(tag=tag, label=kwargs["label"], comp_grid=ctx.obj["compgrid"], ctx=pressure.ctx) - out.push(grid, agyro_vals) - data.add(out) - # end - verb_print(ctx, "Finishing agyro") - - -@click.command() -@click.option("--measure", "-m", default="frobenius", show_default=True, - type=click.Choice(["swidak", "frobenius"]), - help="Specify how to calculate agyrotropy.") -@click.option("--species", "-s", help="Tag for input pressure.") -@click.option("--field", "-f", help="Tag for input EM field.") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.pass_context -def mom_agyro(ctx, **kwargs): - """Compute a measure of agyrotropy. Default measure is taken from - Swisdak 2015. Optionally computes agyrotropy as Frobenius norm of - agyrotropic pressure tensor. - """ - verb_print(ctx, "Starting agyro") - - data = ctx.obj["data"] - tag = "agyro" - if kwargs["tag"]: - tag = kwargs["tag"] - # end - - for species, field in zip(data.iterator(kwargs["species"]), data.iterator(kwargs["field"])): - grid, agyro_vals = get_gkyl_10m_agyro(species=species, field=field, measure=kwargs["measure"]) - out = GData(tag=tag, label=kwargs["label"], comp_grid=ctx.obj["compgrid"], ctx=species.ctx) - out.push(grid, agyro_vals) - data.add(out) - # end - verb_print(ctx, "Finishing agyro") diff --git a/src/postgkyl/commands/animate.py b/src/postgkyl/commands/animate.py deleted file mode 100644 index 54f657da..00000000 --- a/src/postgkyl/commands/animate.py +++ /dev/null @@ -1,493 +0,0 @@ -import os -import shutil -import tempfile -from matplotlib.animation import FuncAnimation, FFMpegWriter -from multiprocessing import Pool -from PIL import Image -import click -import matplotlib -import matplotlib.pyplot as plt -import numpy as np - -from postgkyl.utils import verb_print, set_frame -import postgkyl.output.plot - -# Formats written through ffmpeg (PIL cannot produce these video containers). -VIDEO_EXTS = (".mp4", ".mov", ".avi", ".mkv") - - -def _save_frame_worker(args): - """Worker for parallel frame saving; each process creates its own figure.""" - matplotlib.use("Agg") - frame_idx, frame_data, kwargs, prefix, dpi, figsize = args - fig = plt.figure(figsize=figsize) - _update(0, [frame_data], fig, kwargs) - plt.savefig(f"{prefix:s}_{frame_idx:d}.png", dpi=dpi) - plt.close(fig) -# end - - -def _save_frames(data_list, num_frames, prefix, kwargs, figsize, fig=None): - """Save frames as PNGs, using parallel workers when nproc > 1.""" - if kwargs["nproc"] > 1: - args_list = [(i, data_list[i], kwargs, prefix, kwargs["dpi"], figsize) - for i in range(num_frames)] - with Pool(kwargs["nproc"]) as pool: - pool.map(_save_frame_worker, args_list) - # end - else: - for i in range(num_frames): - _update(i, data_list, fig, kwargs) - plt.savefig(f"{prefix:s}_{i:d}.png", dpi=kwargs["dpi"]) - # end - # end -# end - - -def _compile_movie(frame_files, output_file, fps, duration, ctx): - """Compile PNG frames into an animation.""" - ext = os.path.splitext(output_file)[1].lower() - verb_print(ctx,f"Creating {output_file}...") - if ext in (".gif", ".webp", ".apng"): - images = [Image.open(f) for f in frame_files] - images[0].save( - output_file, save_all=True, append_images=images[1:], - duration=duration, loop=0, optimize=False, - ) - elif ext in VIDEO_EXTS: - # PIL cannot write video containers; use matplotlib's ffmpeg writer. - # duration is in milliseconds per frame, so fall back to it when fps is unset. - movie_fps = fps if fps else 1.0e3 / duration - writer = FFMpegWriter(fps=movie_fps) - first = Image.open(frame_files[0]) - dpi = 100 - fig = plt.figure(figsize=(first.width / dpi, first.height / dpi), dpi=dpi) - ax = fig.add_axes([0, 0, 1, 1]) - ax.axis("off") - with writer.saving(fig, output_file, dpi): - for frame_file in frame_files: - ax.clear() - ax.axis("off") - ax.imshow(Image.open(frame_file)) - writer.grab_frame() - # end - # end - plt.close(fig) - else: - raise ValueError(f"Unsupported output format: {ext}") - - verb_print(ctx,f"{output_file} created.") -# end - - -def _update(frame, data, fig, kwargs): - fig.clear() - kwargs["figure"] = fig - - #global range function is called every frame to set scale limits for frame plot - if kwargs["multiblock"] and kwargs["float"]: - vmin, vmax, num_dims = globalrange(data[frame], kwargs) - if num_dims == 1: - kwargs["ymin"] = vmin - kwargs["ymax"] = vmax - else: - kwargs["zmin"] = vmin - kwargs["zmax"] = vmax - # end - # end - - #main plotting loop - for i, dat in enumerate(data[frame]): - kwargs["title"] = "" - if not kwargs["notitle"]: - if dat.ctx.get("frame") is not None: - kwargs["title"] = f"{kwargs['title']:s} frame: {dat.ctx['frame']:d} " - # end - if dat.ctx.get("time") is not None: - kwargs["title"] = f"{kwargs['title']:s} time: {dat.ctx['time']:.4e}" - # end - # end - - if i == 0: - if kwargs.get("arg"): - im = postgkyl.output.plot(dat, kwargs["arg"], **kwargs) - else: - im = postgkyl.output.plot(dat, **kwargs) - # end - else: - kwargs_ncb = kwargs.copy() - kwargs_ncb["colorbar"] = False - if kwargs.get("arg"): - im = postgkyl.output.plot(dat, kwargs["arg"], **kwargs_ncb) - else: - im = postgkyl.output.plot(dat, **kwargs_ncb) - # end - # end - # end - return im -# end - -#Finds global minima and maxima for all inputed data objects -#also incorporates cutoffglobalrange -def globalrange(data,kwargs): - vmin = float("inf") - vmax = float("-inf") - v_extrema = np.array([]) - for dat in data: - num_dims = dat.get_num_dims() - if num_dims == 1: - val = dat.get_values()*kwargs["yscale"] - else: - val = dat.get_values()*kwargs["zscale"] - # end - if vmin > np.nanmin(val): - vmin = np.nanmin(val) - if vmax < np.nanmax(val): - vmax = np.nanmax(val) - # end - v_extrema = np.append(v_extrema, np.nanmin(val)) - v_extrema = np.append(v_extrema, np.nanmax(val)) - # end - v_extrema = np.sort(v_extrema) - if kwargs["cutoffglobalrange"]: - boundary = 100 * (1 - kwargs["cutoffglobalrange"]) / 2 - vmax = np.percentile(v_extrema, 100 - boundary) - vmin = np.percentile(v_extrema, boundary) - return vmin, vmax, num_dims - else: - return vmin, vmax, num_dims - # end -# end - - -@click.command() -@click.option("--use", "-u", default=None, help="Specify a tag to plot.") -@click.option("--grouptags", is_flag=True, help="Group coresponding tagged frames.") -@click.option("--squeeze", "-p", is_flag=True, help="Squeeze the components into one panel.") -@click.option("--subplots", "-b", is_flag=True, help="Make subplots from multiple datasets.") -@click.option("--nsubplotrow", "nSubplotRow", type=click.INT, - help="Manually set the number of rows for subplots.") -@click.option("--nsubplotcol", "nSubplotCol", type=click.INT, - help="Manually set the number of columns for subplots.") -@click.option("--transpose", is_flag=True, help="Transpose axes.") -@click.option("--contour", "-c", is_flag=True, help="Make contour plot.") -@click.option("--clevels", type=click.STRING, - help="Specify levels for contours: either integer or start:end:nlevels") -@click.option("--quiver", "-q", is_flag=True, help="Make quiver plot.") -@click.option("--streamline", "-l", is_flag=True, help="Make streamline plot.") -@click.option("--sdensity", type=click.FLOAT, help="Control density of the streamlines.") -@click.option("--arrowstyle", type=click.STRING, help="Set the style for streamline arrows.") -@click.option("--group", "-g", type=click.Choice(["0", "1"]), help="Switch to group mode.") -@click.option("--scatter", "-s", is_flag=True, help="Make scatter plot.") -@click.option("--markersize", type=click.FLOAT, help="Set marker size for scatter plots.") -@click.option("--linewidth", type=click.FLOAT, help="Set the linewidth.") -@click.option("--linestyle", type=click.Choice(["solid", "dashed", "dotted", "dashdot"]), - help="Set the linestyle.") -@click.option("--color", type=click.STRING, help="Set color when available.") -@click.option("--style", help="Specify Matplotlib style file (default: Postgkyl).") -@click.option("--diverging", "-d", is_flag=True, help="Switch to diverging colormesh mode.") -@click.option("--arg", type=click.STRING, help="Additional plotting arguments, e.g., '*--'.") -@click.option("--fix-aspect", "-a", "fixaspect", is_flag=True, - help="Enforce the same scaling on both axes.") -@click.option("--logx", is_flag=True, help="Set x-axis to log scale.") -@click.option("--logy", is_flag=True, help="Set y-axis to log scale.") -@click.option("--logz", is_flag=True, help="Set values of 2D plot to log scale.") -@click.option("--xshift", default=0.0, type=click.FLOAT, show_default=True, - help="Value to shift the x-axis.") -@click.option("--yshift", default=0.0, type=click.FLOAT, show_default=True, - help="Value to shift the y-axis.") -@click.option("--zshift", default=0.0, type=click.FLOAT, show_default=True, - help="Value to shift the z-axis.") -@click.option("--xscale", default=1.0, type=click.FLOAT, show_default=True, - help="Value to scale the x-axis.") -@click.option("--yscale", default=1.0, type=click.FLOAT, show_default=True, - help="Value to scale the y-axis.") -@click.option("--zscale", default=1.0, type=click.FLOAT, show_default=True, - help="Value to scale the z-axis.") -@click.option("--float", is_flag=True, - help="Choose min/max levels based on current frame (i.e., each frame uses a different color range).") -@click.option("--xmax", default=None, type=click.FLOAT, help="Set maximal x-value.") -@click.option("--xmin", default=None, type=click.FLOAT, help="Set minimal x-values.") -@click.option("--ymax", default=None, type=click.FLOAT, help="Set maximal y-value.") -@click.option("--ymin", default=None, type=click.FLOAT, help="Set minimal y-values.") -@click.option("--zmax", default=None, type=click.FLOAT, help="Set maximal z-value.") -@click.option("--zmin", default=None, type=click.FLOAT, help="Set minimal z-values.") -@click.option("--xlim", default=None, type=click.STRING, - help="Set limits for the x-coordinate (lower,upper).") -@click.option("--ylim", default=None, type=click.STRING, - help="Set limits for the y-coordinate (lower,upper).") -@click.option("--zlim", default=None, type=click.STRING, - help="Set limits for the z-coordinate (lower,upper).") -@click.option("--cutoffglobalrange", "-cogr", default=None, type=click.FLOAT, - help="Specify middle percentile of data extrema to set y/z limits to") -@click.option("--legend/--no-legend", default=True, help="Show legend.") -@click.option("--colorbar/--no-colorbar", default=True, - help="Show colorbar (2D animations), no colorbar improves animation performance") -@click.option("--force-legend", "forcelegend", is_flag=True, - help="Force legend even when plotting a single dataset.") -@click.option("-x", "--xlabel", type=click.STRING, help="Specify a x-axis label.") -@click.option("-y", "--ylabel", type=click.STRING, help="Specify a y-axis label.") -@click.option("--clabel", type=click.STRING, help="Specify a label for colorbar.") -@click.option("--title", type=click.STRING, help="Specify a title.") -@click.option("--notitle", is_flag=True, help="Do not show title.") -@click.option("-i", "--interval", default=100, help="Specify the animation interval.") -@click.option("--save", is_flag=True, help="Save figure as PNG.") -@click.option("--saveas", type=click.STRING, default=None, help="Name to save the plot as.") -@click.option("--fps", type=click.INT, help="Specify frames per second for saving.") -@click.option("--dpi", type=click.INT, help="DPI (resolution) for output.") -@click.option("--edgecolors", "-e", type=click.STRING, help="Set color for cell edges.") -@click.option("--showgrid/--no-showgrid", default=True, help="Show grid-lines.") -@click.option("--collected", is_flag=True, - help="Animate a dataset that has been collected, i.e. a single dataset with time taken to be the first index.") -@click.option("--hashtag", is_flag=True, help="Turns on the pgkyl hashtag!") -@click.option("--show/--no-show", default=True, help="Turn showing of the plot ON and OFF.") -@click.option("--saveframes", type=click.STRING, - help="Save individual frames as PNGs.") -@click.option("--nproc", default=1, type=click.INT, show_default=True, - help="Number of parallel processes for frame generation.") -@click.option("--tmpdir", default=None, type=click.STRING, show_default=True, - help="Directory to place the temporary directory for parallel frame generation.") -@click.option("--figsize", help="Comma-separated values for x and y size.") -@click.option("-m", "--multiblock", is_flag=True, help="Plots blocks from each frame together") -@click.pass_context -def animate(ctx, **kwargs): - """Animate the actively loaded dataset and show resulting plots in a loop. - - Typically, the datasets are loaded using wildcard/regex feature of the -f option to - the main pgkyl executable. - """ - verb_print(ctx, "Starting animate") - data = ctx.obj["data"] - - # Accept str or path-like input for --saveas (e.g. a pathlib.Path). - if kwargs["saveas"]: - kwargs["saveas"] = str(kwargs["saveas"]) - # end - supported_exts = (".gif", ".webp", ".apng") + VIDEO_EXTS - if kwargs["saveas"] and not kwargs["saveas"].lower().endswith(supported_exts): - raise click.ClickException( - "Unsupported output format for --saveas; please use one of: " - + ", ".join(supported_exts) + ".") - # end - # Video containers are written through ffmpeg, which must be on the PATH. - if kwargs["saveas"] and kwargs["saveas"].lower().endswith(VIDEO_EXTS) \ - and shutil.which("ffmpeg") is None: - raise click.ClickException( - "ffmpeg is required to write " + ", ".join(VIDEO_EXTS) + " files but was " - "not found. Please install ffmpeg or choose a .gif output instead.") - # end - - if kwargs["xlim"]: - kwargs["xmin"] = float(kwargs["xlim"].split(",")[0]) - kwargs["xmax"] = float(kwargs["xlim"].split(",")[1]) - # end - if kwargs["ylim"]: - kwargs["ymin"] = float(kwargs["ylim"].split(",")[0]) - kwargs["ymax"] = float(kwargs["ylim"].split(",")[1]) - # end - if kwargs["zlim"]: - kwargs["zmin"] = float(kwargs["zlim"].split(",")[0]) - kwargs["zmax"] = float(kwargs["zlim"].split(",")[1]) - # end - - if not kwargs["float"] and not kwargs["grouptags"]: - vmin, vmax, num_dims = globalrange(data.iterator(kwargs["use"]), kwargs) - if num_dims == 1: - if kwargs["ymin"] is None: - kwargs["ymin"] = vmin - # end - if kwargs["ymax"] is None: - kwargs["ymax"] = vmax - # end - else: - if kwargs["zmin"] is None: - kwargs["zmin"] = vmin - # end - if kwargs["zmax"] is None: - kwargs["zmax"] = vmax - # end - # end - # end - - anims = [] - figs = [] - kwargs["legend"] = False - - figsize = None - if kwargs["figsize"]: - figsize = (int(kwargs["figsize"].split(",")[0]), int(kwargs["figsize"].split(",")[1])) - # end - - # PIL requires duration in miliseconds. - duration = int(1.0e3 / kwargs["fps"]) if kwargs["fps"] else kwargs["interval"] - - set_figure = False - min_size = np.nan - yset = False - - if kwargs["grouptags"]: - #runs animation for each tag - for tag in data.tag_iterator(kwargs["use"]): - num_datasets = int(data.get_num_datasets(tag=tag)) - min_size = int(np.nanmin((min_size, num_datasets))) - # end - - tag_iterator = list(data.tag_iterator(kwargs["use"])) - kwargs["legend"] = True - set_figure = True - fig_num = int(0) - - for tag in tag_iterator: - #sets scale for each tag animation - vmin, vmax, num_dims = globalrange(data.iterator(tag), kwargs) - if num_dims == 1: - kwargs["ymin"] = vmin - kwargs["ymax"] = vmax - yset = True - else: - if yset: #so that ymin,ymax of 1D anim don't affect 2D anim - kwargs["ymin"] = None - kwargs["ymax"] = None - # end - kwargs["zmin"] = vmin - kwargs["zmax"] = vmax - # end - - #creating min list of lists (non-multiblock case) - data_list = [] - for dat in data.iterator(tag): - data_list.append([dat]) - # end - figs.append(plt.figure(fig_num, figsize=figsize)) - fig_num += 1 - - num_frames = int(np.nanmin((min_size, len(data_list)))) - file_name = f"anim_{tag:s}.gif" if tag is not None else "anim.gif" - if kwargs["saveas"]: - file_name = str(kwargs["saveas"]) - # end - - if kwargs["saveframes"]: - # Save PNGs, then optionally compile a movie. - _save_frames(data_list, num_frames, kwargs["saveframes"], kwargs, figsize, figs[-1]) - if kwargs["save"] or kwargs["saveas"]: - frame_files = [f"{kwargs['saveframes']}_{i}.png" for i in range(num_frames)] - _compile_movie(frame_files, file_name, kwargs["fps"], duration, ctx) - # end - kwargs["show"] = False - elif kwargs["nproc"] > 1: - # Parallel: use a temp dir, compile, then clean up. - with tempfile.TemporaryDirectory(dir=kwargs["tmpdir"]) as tmpdir: - tmp_prefix = os.path.join(tmpdir, "frame") - _save_frames(data_list, num_frames, tmp_prefix, kwargs, figsize) - frame_files = [f"{tmp_prefix}_{i}.png" for i in range(num_frames)] - _compile_movie(frame_files, file_name, kwargs["fps"], duration, ctx) - # end - kwargs["show"] = False - else: - anims.append( - FuncAnimation(figs[-1], _update, num_frames, - fargs=(data_list, figs[-1], kwargs), interval=kwargs["interval"], - blit=False) - ) - if kwargs["save"] or kwargs["saveas"]: - anims[-1].save(file_name, writer="ffmpeg", fps=kwargs["fps"], dpi=kwargs["dpi"]) - # end - # end - # end - #animation code for multiblock case - elif kwargs["multiblock"]: - - #set ctx frames for all data objects - sorted_frame_list = set_frame(ctx) - - #create main list of lists (multiblock case) - data_list = [] - #organize data objects so each interior list includes blocks from one frame - for frame in sorted_frame_list: - frame_data_list = [dat for dat in data.iterator(kwargs["use"]) if dat.ctx["frame"] == frame] - data_list.append(frame_data_list) - # end - - figs.append(plt.figure(figsize=figsize)) - #makes default color blue in 1D cases, this prevents blocks from having different colors - if (not kwargs["color"] and data_list[0][0].get_num_dims() == 1): - kwargs["color"] = "tab:blue" - # end - - num_frames = int(np.nanmin((min_size, len(data_list)))) - file_name = kwargs["saveas"] if kwargs["saveas"] else "anim.gif" - - if kwargs["saveframes"]: - _save_frames(data_list, num_frames, kwargs["saveframes"], kwargs, figsize, figs[-1]) - if kwargs["save"] or kwargs["saveas"]: - frame_files = [f"{kwargs['saveframes']}_{i}.png" for i in range(num_frames)] - _compile_movie(frame_files, file_name, kwargs["fps"], duration, ctx) - # end - kwargs["show"] = False - elif kwargs["nproc"] > 1: - with tempfile.TemporaryDirectory(dir=kwargs["tmpdir"]) as tmpdir: - tmp_prefix = os.path.join(tmpdir, "frame") - _save_frames(data_list, num_frames, tmp_prefix, kwargs, figsize) - frame_files = [f"{tmp_prefix}_{i}.png" for i in range(num_frames)] - _compile_movie(frame_files, file_name, kwargs["fps"], duration, ctx) - # end - kwargs["show"] = False - else: - anims.append( - FuncAnimation(figs[-1], _update, num_frames, - fargs=(data_list, figs[-1], kwargs), interval=kwargs["interval"], - blit=False) - ) - if kwargs["save"] or kwargs["saveas"]: - anims[-1].save(file_name, writer="ffmpeg", fps=kwargs["fps"], dpi=kwargs["dpi"]) - # end - # end - - else: - - #create main list of lists (non-multiblock case) - data_list = [] - for dat in data.iterator(kwargs["use"]): - data_list.append([dat]) - # end - if set_figure: - figs.append(plt.figure(fig_num, figsize=figsize)) - else: - figs.append(plt.figure(figsize=figsize)) - # end - - num_frames = int(np.nanmin((min_size, len(data_list)))) - file_name = kwargs["saveas"] if kwargs["saveas"] else "anim.gif" - - if kwargs["saveframes"]: - _save_frames(data_list, num_frames, kwargs["saveframes"], kwargs, figsize, figs[-1]) - if kwargs["save"] or kwargs["saveas"]: - frame_files = [f"{kwargs['saveframes']}_{i}.png" for i in range(num_frames)] - _compile_movie(frame_files, file_name, kwargs["fps"], duration, ctx) - # end - kwargs["show"] = False - elif kwargs["nproc"] > 1: - with tempfile.TemporaryDirectory(dir=kwargs["tmpdir"]) as tmpdir: - tmp_prefix = os.path.join(tmpdir, "frame") - _save_frames(data_list, num_frames, tmp_prefix, kwargs, figsize) - frame_files = [f"{tmp_prefix}_{i}.png" for i in range(num_frames)] - _compile_movie(frame_files, file_name, kwargs["fps"], duration, ctx) - # end - kwargs["show"] = False - else: - anims.append( - FuncAnimation(figs[-1], _update, num_frames, - fargs=(data_list, figs[-1], kwargs), interval=kwargs["interval"], - blit=False) - ) - if kwargs["save"] or kwargs["saveas"]: - anims[-1].save(file_name, writer="ffmpeg", fps=kwargs["fps"], dpi=kwargs["dpi"]) - # end - # end - # end - - if kwargs["show"]: - plt.show() - # end - verb_print(ctx, "Finishing animate") diff --git a/src/postgkyl/commands/bparrotate.py b/src/postgkyl/commands/bparrotate.py deleted file mode 100644 index 1f4a39e8..00000000 --- a/src/postgkyl/commands/bparrotate.py +++ /dev/null @@ -1,44 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.parrotate - - -@click.command() -@click.option("--array", "-a", default="array", show_default=True, - help="Tag for array to be rotated") -@click.option("--field", "-r", default="field", show_default=True, - help="Tag for EM field data (data used for the rotation)") -@click.option("--tag", "-t", default="arrayBpar", show_default=True, - help="Tag for the resulting rotated array parallel to magnetic field") -@click.option("--label", "-l", default="arrayBpar", show_default=True, - help="Custom label for the result") -@click.pass_context -def bparrotate(ctx, **kwargs): - """Rotate an array parallel to the unit vectors of the magnetic field. - - For two arrays u and b, where b is the unit vector in the direction of the magnetic - field, the operation is (u dot b_hat) b_hat. Note that the magnetic field is a - three-component field, so the output is a new vector whose components are (u_{b_x}, - u_{b_y}, u_{b_z}), i.e., the x, y, and z components of the vector u parallel to the - magnetic field. - """ - verb_print(ctx, "Starting rotation parallel to magnetic field") - - data = ctx.obj["data"] # shortcut - - for a, rot in zip(data.iterator(kwargs["array"]), data.iterator(kwargs["field"])): - # Magnetic field is components 3, 4, & 5 in field array - grid, outrot = postgkyl.tools.parrotate(a, rot, "3:6") - # Create new GData structure with appropriate outtag and labels to store output. - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=a.ctx) - out.push(grid, outrot) - data.add(out) - # end - - data.deactivate_all(tag=kwargs["array"]) - data.deactivate_all(tag=kwargs["field"]) - - verb_print(ctx, "Finishing rotation parallel to magnetic field") diff --git a/src/postgkyl/commands/bperprotate.py b/src/postgkyl/commands/bperprotate.py deleted file mode 100644 index 630ff97e..00000000 --- a/src/postgkyl/commands/bperprotate.py +++ /dev/null @@ -1,41 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.perprotate - - -@click.command() -@click.option("--array", "-a", default="array", show_default=True, - help="Tag for array to be rotated.") -@click.option("--field", "-r", default="field", show_default=True, - help="Tag for EM field data (data used for the rotation).") -@click.option("--tag", "-t", default="arrayBperp", show_default=True, - help="Tag for the resulting rotated array perpendicular to magnetic field.") -@click.option("--label", "-l", default="arrayBperp", show_default=True, - help="Custom label for the result.") -@click.pass_context -def bperprotate(ctx, **kwargs): - """Rotate an array perpendicular to the unit vectors of the magnetic field. - - For two arrays u and b, where b is the unit vector in the direction of the magnetic - field, the operation is u - (u dot b_hat) b_hat. - """ - verb_print(ctx, "Starting rotation perpendicular to magnetic field") - - data = ctx.obj["data"] # shortcut - - for a, rot in zip(data.iterator(kwargs["array"]), data.iterator(kwargs["field"])): - # Magnetic field is components 3, 4, & 5 in field array - grid, outrot = postgkyl.tools.perprotate(a, rot, "3:6") - # Create new GData structure with appropriate outtag and labels to store output. - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=a.ctx) - out.push(grid, outrot) - data.add(out) - # end - - data.deactivate_all(tag=kwargs["array"]) - data.deactivate_all(tag=kwargs["field"]) - - verb_print(ctx, "Finishing rotation perpendicular to magnetic field") diff --git a/src/postgkyl/commands/collect.py b/src/postgkyl/commands/collect.py deleted file mode 100644 index 5fc3509c..00000000 --- a/src/postgkyl/commands/collect.py +++ /dev/null @@ -1,115 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print - - -@click.command() -@click.option("-s", "--sumdata", is_flag=True, - help="Sum data in the collected datasets (retain components).") -@click.option("-p", "--period", type=click.FLOAT, - help="Specify a period to create epoch data instead of time data.") -@click.option("--offset", default=0.0, type=click.FLOAT, show_default=True, - help="Specify an offset to create epoch data instead of time data.") -@click.option("-c", "--chunk", type=click.INT, - help="Collect into chunks with specified length rather than into a single dataset.") -@click.option("--use", "-u", default=None, help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", default=None, help="Specify a 'tag' for the result.") -@click.option("--label", "-l", default=None, help="Specify the custom label for the result.") -@click.pass_context -def collect(ctx, **kwargs): - """Collect data from the active datasets and create a new combined dataset. - - The time-stamp in each of the active datasets is collected and used as the new X-axis. - Data can be collected in chunks, in which case several datasets are created, each with - the chunk-sized pieces collected into each new dataset. - """ - verb_print(ctx, "Starting collect") - data = ctx.obj["data"] - - if kwargs["tag"]: - out_tags = kwargs["tag"].split(",") - # end - - tag_cnt = 0 - for tag in data.tag_iterator(kwargs["use"]): - time = [[]] - values = [[]] - grid = [[]] - cnt = 0 - label = None - - for i, dat in data.iterator(tag, enum=True): - cnt += 1 - if kwargs["chunk"] and cnt > kwargs["chunk"]: - cnt = 1 - time.append([]) - values.append([]) - grid.append([]) - # end - if dat.ctx["time"]: - time[-1].append(dat.ctx["time"]) - elif dat.ctx["frame"]: - time[-1].append(dat.ctx["frame"]) - else: - time[-1].append(i) - # end - val = dat.get_values() - if kwargs["sumdata"]: - num_dims = dat.get_num_dims() - axis = tuple(range(num_dims)) - values[-1].append(np.nansum(val, axis=axis)) - else: - values[-1].append(val) - # end - if not grid[-1]: - grid[-1] = dat.get_grid().copy() - # end - label = dat.get_custom_label() - # end - - data.deactivate_all(tag) - - out_tag = tag - if kwargs["tag"]: - if len(out_tags) > 1: - out_tag = out_tags[tag_cnt] - else: - out_tag = out_tags[0] - # end - # end - tag_cnt += 1 - - if label is None: - label = "collect" - # end - if kwargs["label"]: - label = kwargs["label"] - # end - - for i in range(len(time)): - time[i] = np.array(time[i]) - values[i] = np.array(values[i]) - - if kwargs.get("period"): - time[i] = (time[i] - kwargs["offset"]) % kwargs["period"] - # end - - sort_idx = np.argsort(time[i]) - time[i] = time[i][sort_idx] - values[i] = values[i][sort_idx] - - if kwargs["sumdata"]: - grid[i] = [time[i]] - else: - grid[i].insert(0, np.array(time[i])) - # end - - out = GData(tag=out_tag, label=label, comp_grid=ctx.obj["compgrid"]) - out.push(grid[i], values[i]) - data.add(out) - # end - # end - - verb_print(ctx, "Finishing collect") diff --git a/src/postgkyl/commands/config.py b/src/postgkyl/commands/config.py deleted file mode 100644 index fb1a2450..00000000 --- a/src/postgkyl/commands/config.py +++ /dev/null @@ -1,27 +0,0 @@ -import os -import pathlib - -import click - -from postgkyl._gkylsoft_path import default_config_path - -@click.command(name="config") -@click.option("--gkylsoft", "-g", default=None, type=click.Path(), - help="Path to the gkylsoft directory. Uses GKYLSOFT_DIR env variable if not provided.") -@click.option("--config-file", "-c", default=None, type=click.Path(), - help="Config file to write. Default: ~/.postgkyl/gkylsoft_path, " - "or the POSTGKYL_CONFIG env variable if set.") -def config(gkylsoft, config_file): - """Write postgkyl configuration (gkylsoft path) to the config file.""" - - if gkylsoft is None: - gkylsoft = os.environ.get("GKYLSOFT_DIR") - - if gkylsoft is None: - raise click.UsageError("No gkylsoft path provided. Pass --gkylsoft /path/to/gkylsoft " - "or set the GKYLSOFT_DIR env variable.") - - out = pathlib.Path(config_file if config_file is not None else default_config_path()) - out.parent.mkdir(parents=True, exist_ok=True) - out.write_text(f"GKYLSOFT_DIR={gkylsoft}\n") - click.echo(f"Wrote gkylsoft path to {out}") diff --git a/src/postgkyl/commands/current.py b/src/postgkyl/commands/current.py deleted file mode 100644 index 6563b2b6..00000000 --- a/src/postgkyl/commands/current.py +++ /dev/null @@ -1,32 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.accumulate_current - - -@click.command() -@click.option("--qbym", "-q", default=False, show_default=True, - help="Flag for multiplying by charge/mass ratio instead of just charge.") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", default="current", show_default=True, - help="Tag for the resulting current array.") -@click.option("--label", "-l", default="J", show_default=True, help="Custom label for the result.") -@click.pass_context -def current(ctx, **kwargs): - """Accumulate current, sum over species of charge multiplied by flow.""" - verb_print(ctx, "Starting current accumulation") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - grid = dat.get_grid() - outcurrent = np.zeros(dat.get_values().shape) - grid, outcurrent = postgkyl.tools.accumulate_current(dat, kwargs["qbym"]) - dat.deactivate() - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=dat.ctx) - out.push(grid, outcurrent) - data.add(out) - # end - verb_print(ctx, "Finishing current accumulation") diff --git a/src/postgkyl/commands/data_space.py b/src/postgkyl/commands/data_space.py deleted file mode 100644 index 1d2a7b00..00000000 --- a/src/postgkyl/commands/data_space.py +++ /dev/null @@ -1,179 +0,0 @@ -"""Postgkyl submodule to provide iterators in hte command line mode.""" -from __future__ import annotations - -import click -import numpy as np -from typing import Iterator, TYPE_CHECKING - -if TYPE_CHECKING: - from postgkyl import GData -#end - -class DataSpace(object): - """Postgkyl class to store information about datasets and provide iterators in the command line mode.""" - - def __init__(self): - self._dataset_dict = {} - - # ---- Iterators ---- - def iterator(self, tag: str | None = None, enum: bool = False, - only_active: bool = True, select: int | slice | str | None = None) -> Iterator[GData]: - # Process 'select' - if enum and select: - click.echo(click.style("Error: 'select' and 'enum' cannot be selected simultaneously", fg="red")) - quit() - # end - idx_sel = slice(None, None) - if isinstance(select, int): - idx_sel = [select] - elif isinstance(select, slice): - idx_sel = select - elif isinstance(select, str): - if ":" in select: - lo = None - up = None - step = None - s = select.split(":") - if s[0]: - lo = int(s[0]) - # end - if s[1]: - up = int(s[1]) - # end - if len(s) > 2: - step = int(s[2]) - # end - idx_sel = slice(lo, up, step) - else: - idx_sel = list([int(s) for s in select.split(",")]) - # end - # end - - if tag: - tags = tag.split(",") - else: - tags = list(self._dataset_dict) - # end - for t in tags: - try: - if not select or isinstance(idx_sel, slice): - for i, dat in enumerate(self._dataset_dict[t][idx_sel]): - if (not only_active) or dat.get_status(): # implication - if enum: - yield i, dat - else: - yield dat - # end - # end - # end - else: # isinstance(idx_sel, list) - for i in idx_sel: - dat = self._dataset_dict[t][i] - if (not only_active) or dat.get_status(): # implication - yield dat - # end - # end - # end - except KeyError as err: - click.echo(click.style(f"ERROR: Failed to load the specified/default tag {err}", fg="red")) - quit() - except IndexError: - click.echo(click.style("ERROR: Index out of the dataset range", fg="red")) - quit() - # end - # end - - def tag_iterator(self, tag: str | None = None, only_active: bool = True) -> Iterator[str]: - if tag: - out = tag.split(",") - elif only_active: - out = [] - for t in self._dataset_dict: - if True in (dat.get_status() for dat in self.iterator(t)): - out.append(t) - # end - # end - else: - out = list(self._dataset_dict) - # end - return iter(out) - - # ---- Labels ---- - def set_unique_labels(self) -> None: - num_comps = [] - names = [] - labels = [] - for dat in self.iterator(): - file_name = dat._file_name - extension_len = len(file_name.split(".")[-1]) - file_name = file_name[: -(extension_len + 1)] - # only remove the file extension but take into account - # that the file name might start with '../' - sp = file_name.split("_") - names.append(sp) - num_comps.append(int(len(sp))) - labels.append("") - # end - max_elem = np.max(num_comps) - idx_max = np.argmax(num_comps) - for i in range(max_elem): - include = False - reference = names[idx_max][i] - for nm in names: - if i < len(nm) and nm[i] != reference: - include = True - # end - # end - if include: - for idx, nm in enumerate(names): - if i < len(nm): - if labels[idx] == "": - labels[idx] += nm[i] - else: - labels[idx] += f"_{nm[i]:s}" - # end - # end - # end - # end - # end - cnt = 0 - for dat in self.iterator(): - dat.set_label(labels[cnt]) - cnt += 1 - # end - - # ---- Adding datasets ---- - def add(self, data: GData) -> None: - tag_nm = data.get_tag() - if tag_nm in self._dataset_dict: - self._dataset_dict[tag_nm].append(data) - else: - self._dataset_dict[tag_nm] = [data] - # end - - # ---- Staus control ---- - def activate_all(self, tag: str | None = None) -> None: - for dat in self.iterator(tag=tag, only_active=False): - dat.deactivate() - # end - - # end - def deactivate_all(self, tag: str | None = None) -> None: - for dat in self.iterator(tag=tag, only_active=False): - dat.deactivate() - # end - - # ---- Utilities ---- - def get_dataset(self, idx: int, tag: str = "default") -> GData: - return self._dataset_dict[tag][idx] - - - def get_num_datasets(self, tag: str | None = None, only_active: bool = True): - num_sets = 0 - for dat in self.iterator(tag=tag, only_active=only_active): - num_sets += 1 - # end - return num_sets - - def clean(self): - self._dataset_dict = {} \ No newline at end of file diff --git a/src/postgkyl/commands/dg_avg.py b/src/postgkyl/commands/dg_avg.py deleted file mode 100644 index 6533210c..00000000 --- a/src/postgkyl/commands/dg_avg.py +++ /dev/null @@ -1,101 +0,0 @@ -import os - -import click - -from postgkyl.data import GData -from postgkyl.data import select as pgkyl_select -from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops -from postgkyl.utils import verb_print - - -def _jacobgeo_path(file_name): - """Return the -geo_int_jacobgeo.gkyl path next to file_name, or None.""" - if not file_name: - return None - dirname = os.path.dirname(file_name) - basename = os.path.basename(file_name) - prefix = basename.split("-")[0] - return os.path.join(dirname, f"{prefix}-geo_int_jacobgeo.gkyl") - - -@click.command(name="dg-avg") -@click.option("--z0", is_flag=True, help="Average over direction 0.") -@click.option("--z1", is_flag=True, help="Average over direction 1.") -@click.option("--z2", is_flag=True, help="Average over direction 2.") -@click.option("--z3", is_flag=True, help="Average over direction 3.") -@click.option("--z4", is_flag=True, help="Average over direction 4.") -@click.option("--z5", is_flag=True, help="Average over direction 5.") -@click.option("--comp", "-c", default=None, - help="Component index to select from the result (int or slice).") -@click.option("--weight", "-w", default=None, - help="Weight file for the average. Defaults to -geo_int_jacobgeo.gkyl " - "found next to the dataset; pass a path to override, or 'none' to disable.") -@click.option("--use", "-u", help="Tag to apply to. [default: all active]") -@click.option("--tag", "-t", help="Tag for the output dataset.") -@click.option("--label", "-l", help="Label for the output dataset.") -@click.pass_context -def dg_avg(ctx, **kwargs): - """ - Average a DG field over specified directions. - - Directions to average over are specified using the flags --z0, --z1, ... --z5. - The output has a reduced dimensionality corresponding to the averaged directions. - - The geometric Jacobian -geo_int_jacobgeo.gkyl (found next to the dataset) - is used as the weight if present, giving the weighted average int(f J dx) / int(J dx). - Override the weight file with --weight, or disable weighting with '--weight none'. - """ - verb_print(ctx, "Starting dg-avg") - data = ctx.obj["data"] - - z_opts = [kwargs["z0"], kwargs["z1"], kwargs["z2"], - kwargs["z3"], kwargs["z4"], kwargs["z5"]] - avg_dirs = [i for i, z in enumerate(z_opts) if z] - - if not avg_dirs: - ctx.fail("dg-avg requires at least one direction flag (--z0 ... --z5).") - - ops = GkeyllDGops() - - weight_opt = kwargs["weight"] - weight_off = weight_opt is not None and str(weight_opt).strip().lower() == "none" - weight_cache = {} # file path -> loaded GData (avoid re-reading per dataset) - - def _load_weight(dat): - """Resolve and load the weight GData for a dataset, or None.""" - if weight_off: - return None - if weight_opt: # explicit override: must exist - path = weight_opt - if not os.path.isfile(path): - ctx.fail(f"weight file '{path}' not found.") - else: # auto-detect the geometric Jacobian - path = _jacobgeo_path(dat.get_file_name()) - if path is None or not os.path.isfile(path): - verb_print(ctx, f"No jacobgeo weight found ({path}); using unweighted average.") - return None - if path not in weight_cache: - verb_print(ctx, f"Loading average weight from {path}") - weight_cache[path] = GData(file_name=path, comp_grid=ctx.obj["compgrid"]) - return weight_cache[path] - - for dat in data.iterator(kwargs["use"]): - out = dat - weight = _load_weight(dat) - - # Perform the averaging directly on the specified directions - out = ops.average(avg_dirs, out, weight=weight, - comp_grid=ctx.obj["compgrid"]) - - if kwargs["tag"]: - out.set_tag(kwargs["tag"]) - if kwargs["label"]: - out.set_label(kwargs["label"]) - - if kwargs["comp"] is not None: - pgkyl_select(out, overwrite=True, comp=kwargs["comp"]) - - dat.deactivate() - data.add(out) - - verb_print(ctx, "Finishing dg-avg") \ No newline at end of file diff --git a/src/postgkyl/commands/dg_evproj.py b/src/postgkyl/commands/dg_evproj.py deleted file mode 100644 index fce9461b..00000000 --- a/src/postgkyl/commands/dg_evproj.py +++ /dev/null @@ -1,59 +0,0 @@ -import click - -from postgkyl.data import select as pgkyl_select -from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops -from postgkyl.utils import verb_print - -@click.command(name="dg-evproj") -@click.option("--z0", default=None, type=float, - help="Physical coord to evaluate in direction 0.") -@click.option("--z1", default=None, type=float, - help="Physical coord to evaluate in direction 1.") -@click.option("--z2", default=None, type=float, - help="Physical coord to evaluate in direction 2.") -@click.option("--z3", default=None, type=float, - help="Physical coord to evaluate in direction 3.") -@click.option("--z4", default=None, type=float, - help="Physical coord to evaluate in direction 4.") -@click.option("--z5", default=None, type=float, - help="Physical coord to evaluate in direction 5.") -@click.option("--comp", "-c", default=None, - help="Component index to select from the result (int or slice).") -@click.option("--use", "-u", help="Tag to apply to. [default: all active]") -@click.option("--tag", "-t", help="Tag for the output dataset.") -@click.option("--label", "-l", help="Label for the output dataset.") -@click.pass_context -def dg_evproj(ctx, **kwargs): - """ - Evaluate a DG field at specified coordinates and project onto a lower-dimensional basis. - - Coordinates specified with --z0, --z1, ... --z5. - """ - verb_print(ctx, "Starting dg-evproj") - data = ctx.obj["data"] - - z_opts = [kwargs["z0"], kwargs["z1"], kwargs["z2"], - kwargs["z3"], kwargs["z4"], kwargs["z5"]] - eval_dirs = [i for i, z in enumerate(z_opts) if z is not None] - eval_coords = [z_opts[i] for i in eval_dirs] - - if not eval_dirs: - ctx.fail("dg-evproj requires at least one --z0 ... --z5 coordinate.") - - ops = GkeyllDGops() - - for dat in data.iterator(kwargs["use"]): - out = ops.eval_at_coord_proj(eval_dirs, eval_coords, dat, - comp_grid=ctx.obj["compgrid"]) - if kwargs["tag"]: - out.set_tag(kwargs["tag"]) - if kwargs["label"]: - out.set_label(kwargs["label"]) - - if kwargs["comp"] is not None: - pgkyl_select(out, overwrite=True, comp=kwargs["comp"]) - - dat.deactivate() - - data.add(out) - verb_print(ctx, "Finishing dg-evproj") diff --git a/src/postgkyl/commands/dg_local_poly.py b/src/postgkyl/commands/dg_local_poly.py deleted file mode 100644 index 303277ef..00000000 --- a/src/postgkyl/commands/dg_local_poly.py +++ /dev/null @@ -1,135 +0,0 @@ -import click -import numpy as np - -from postgkyl.utils import verb_print -from postgkyl.data.dg import _getnum_nodes -from postgkyl.modalDG.kernels import expand_1d, expand_2d, expand_3d, expand_4d, expand_5d, expand_6d - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--npoints", "-n", type=click.INT, default=2, - help="Number of evaluation points per cell.") -@click.pass_context -def dg_local_poly(ctx, **kwargs): - """ - Generate a discontinuous DG polynomial cellwise representation of the data. - The modal DG decomposition is evaluated with npoints per cell from one face - to the other. A NaN is inserted at every cell interface so that, when plotted, - the curve is broken at each interface and the inter-cell discontinuities of the DG solution - are visible. - Example (1D plot of the M0 moment along x at frame 0): - pgkyl sim_3x2v_p1-ion_M0_0.gkyl dg-local-poly sel --z1=0.0 --z2=0.0 pl - """ - verb_print(ctx, "Starting dg-local-poly") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - poly_order = dat.ctx.get("poly_order") - - if poly_order is None: - ctx.fail(click.style( - "ERROR in dg-local-poly: no 'poly_order' was specified and dataset " - f"{dat.get_label():s} does not have the required information.", - fg="red")) - - num_dims = dat.get_num_dims() - - num_cells = dat.get_num_cells() - values = dat.get_values() - - num_basis = int(_getnum_nodes(num_dims, poly_order, "serendipity")) - num_eqn = int(dat.get_num_comps() // num_basis) - - # Reference evaluation nodes: just inside the two cell interfaces. - nodes = np.linspace(-1.0, 1.0, kwargs["npoints"]) - num_nodes = len(nodes) - - # Evaluate the modal decomposition of each field at the interface nodes. - int_values = np.zeros(tuple(np.int32(num_cells * num_nodes)) + (num_eqn,)) - for m in range(num_eqn): - # Raw modal coefficients of field m, shape (..., num_basis). - q = values[..., m * num_basis:(m + 1) * num_basis] - if num_dims == 1: - for i, x in enumerate(nodes): - int_values[i::num_nodes, m] = expand_1d[int(poly_order - 1)](q, x) - elif num_dims == 2: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - int_values[i::num_nodes, j::num_nodes, m] = expand_2d[ - int(poly_order - 1)](q, x, y) - elif num_dims == 3: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - int_values[i::num_nodes, j::num_nodes, k::num_nodes, m] = expand_3d[ - int(poly_order - 1)](q, x, y, z) - elif num_dims == 4: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - for l, v1 in enumerate(nodes): - int_values[i::num_nodes, j::num_nodes, k::num_nodes, l::num_nodes, - m] = expand_4d[int(poly_order - 1)](q, x, y, z, v1) - elif num_dims == 5: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - for l, v1 in enumerate(nodes): - for m1, v2 in enumerate(nodes): - int_values[i::num_nodes, j::num_nodes, k::num_nodes, - l::num_nodes, m1::num_nodes, m] = expand_5d[ - int(poly_order - 1)](q, x, y, z, v1, v2) - elif num_dims == 6: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - for l, v1 in enumerate(nodes): - for m1, v2 in enumerate(nodes): - for n1, v3 in enumerate(nodes): - int_values[i::num_nodes, j::num_nodes, k::num_nodes, - l::num_nodes, m1::num_nodes, n1::num_nodes, - m] = expand_6d[int(poly_order - 1)](q, x, y, z, v1, - v2, v3) - # Build the grid with the physical coordinates of the nodes. - grid_in = dat.get_grid() - lower, upper = dat.get_bounds() - int_grid = [] - for d in range(num_dims): - g = np.squeeze(np.asarray(grid_in[d])) - if g.ndim == 1 and g.shape[0] == num_cells[d] + 1: - edges_d = g - else: - edges_d = np.linspace(lower[d], upper[d], num_cells[d] + 1) - cell_center = 0.5 * (edges_d[:-1] + edges_d[1:]) - dx = edges_d[1:] - edges_d[:-1] - coords = (cell_center[:, np.newaxis] + nodes[np.newaxis, :]*dx[:, np.newaxis]/2).reshape(-1) - int_grid.append(coords) - - if dat.ctx["grid_type"] == "c2p_vel": - # Evaluate the vel map at nodes in each cell. - poly_order_vmap = 1 - num_cdim = dat.ctx["num_cdim"] - num_vdim = dat.ctx["num_vdim"] - for d in range(num_cdim,num_dims): - grid_c = grid_in[d] - coord_1v = np.zeros(num_cells[d]*num_nodes) - for i, vmap_c in enumerate(grid_c): - for k in range(num_nodes): - coord_1v[i*num_nodes+k] = expand_1d[int(poly_order_vmap - 1)](vmap_c, nodes[k]) - # end - # end - - int_grid[d] = coord_1v - # end - # end - - # Insert a NaN between every couple of points along each dimension to break - # the curve at the cell interfaces. - for d in range(num_dims): - sep = np.arange(num_nodes, num_nodes * num_cells[d], num_nodes) - int_values = np.insert(int_values, sep, np.nan, axis=d) - int_grid[d] = np.insert(int_grid[d], sep, int_grid[d][sep - 1]) - - dat.push(int_grid, int_values) - verb_print(ctx, "Finishing dg-local-poly") diff --git a/src/postgkyl/commands/differentiate.py b/src/postgkyl/commands/differentiate.py deleted file mode 100644 index 0edb0e13..00000000 --- a/src/postgkyl/commands/differentiate.py +++ /dev/null @@ -1,67 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.data import GInterpModal, GInterpNodal -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--basis_type", "-b", type=click.Choice(["ms", "ns", "mo"]), help="Specify DG basis.") -@click.option("--poly_order", "-p", type=click.INT, help="Specify polynomial order.") -@click.option("--interp", "-i", type=click.INT, - help="Interpolation onto a general mesh of specified amount") -@click.option("--direction", "-d", type=click.INT, - help="Direction of the derivative. [default: calculate all]") -@click.option("--read", "-r", type=click.BOOL, help="Read from general interpolation file.") -@click.option("--use", "-u", help="Specify a 'tag' to apply to. [default: all]") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def differentiate(ctx, **kwargs): - """Interpolate a derivative of DG data on a uniform mesh.""" - verb_print(ctx, "Starting differentiate") - data = ctx.obj["data"] - - basis_type = None - is_modal = None - if kwargs.get("basis_type"): - if kwargs["basis_type"] == "ms": - basis_type = "serendipity" - is_modal = True - elif kwargs["basis_type"] == "ns": - basis_type = "serendipity" - is_modal = False - elif kwargs["basis_type"] == "mo": - basis_type = "maximal-order" - is_modal = True - elif kwargs["basis_type"] == "mt": - basis_type = "tensor" - is_modal = True - # end - # end - - for dat in data.iterator(kwargs["use"]): - if kwargs["basis_type"] is None and dat.ctx["basis_type"] is None: - ctx.fail( - click.style(f"ERROR in interpolate: no 'basis_type' was specified and dataset {dat.get_label():s} does not have required ctxdata", - fg="red") - ) - # end - - if is_modal or dat.ctx["is_modal"]: - dg = GInterpModal(dat, kwargs["poly_order"], kwargs["basis_type"], - kwargs["interp"], kwargs["read"]) - else: - dg = GInterpNodal(dat, kwargs["poly_order"], basis_type, kwargs["interp"], kwargs["read"]) - # end - - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - grid, values = dg.differentiate(direction=kwargs["direction"]) - out.push(grid, values) - data.add(out) - else: - dg.differentiate(direction=kwargs["direction"], overwrite=True) - # end - verb_print(ctx, "Finishing differentiate") diff --git a/src/postgkyl/commands/energetics.py b/src/postgkyl/commands/energetics.py deleted file mode 100644 index dd033238..00000000 --- a/src/postgkyl/commands/energetics.py +++ /dev/null @@ -1,37 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -import postgkyl.tools.energetics - - -@click.command() -@click.option("--elc", "-e", default="elc", show_default=True, help="Tag for electrons.") -@click.option("--ion", "-i", default="ion", show_default=True, help="Tag for ions.") -@click.option("--field", "-f", default="field", show_default=True, help="Tag for EM fields.") -@click.option("--tag", "-t", default="energetics", show_default=True, help="Tag for the result.") -@click.option("--label", "-l", default="E", show_default=True, help="Custom label for the result.") -@click.pass_context -def energetics(ctx, **kwargs): - """Decomposes the components of the energy (kinetic, thermal, electromagnetic) for a two-species (electron, ion) plasma.""" - verb_print(ctx, "Starting energetics decomposition") - data = ctx.obj["data"] # shortcut - - for elc, ion, em in zip(data.iterator(kwargs["elc"]), - data.iterator(kwargs["ion"]), data.iterator(kwargs["field"])): - grid = em.get_grid() - out_energetics = np.zeros(em.get_values()[..., 0:7].shape) - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=em.ctx) - grid, out_energetics = postgkyl.tools.energetics(elc, ion, em) - out.push(grid, out_energetics) - data.add(out) - # end - - data.deactivate_all(tag=kwargs["elc"]) - data.deactivate_all(tag=kwargs["ion"]) - data.deactivate_all(tag=kwargs["field"]) - - verb_print(ctx, "Finishing energetics decomposition") diff --git a/src/postgkyl/commands/euler.py b/src/postgkyl/commands/euler.py deleted file mode 100644 index 301eb7be..00000000 --- a/src/postgkyl/commands/euler.py +++ /dev/null @@ -1,57 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -import postgkyl.tools.prim_vars as pv - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("-g", "--gas_gamma", type=click.FLOAT, default=5.0/3.0, show_default=True, - help="Gas adiabatic constant.") -@click.option("-v", "--variable_name", prompt=True, - type=click.Choice(["density", "xvel", "yvel", "zvel", "vel", "pressure", "ke", "temp", "sound", "mach"]), - help="Variable to extract.") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def euler(ctx, **kwargs): - """Compute Euler (five-moment) primitive and some derived variables - from fluid conserved variables. - """ - verb_print(ctx, "Starting euler") - data = ctx.obj["data"] - - v = kwargs["variable_name"] - for dat in data.iterator(kwargs["use"]): - verb_print(ctx, f"euler: Extracting {v:s} from data set.") - out = dat - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - data.add(out) - # end - if v == "density": - pv.get_density(dat, out_mom=out) - elif v == "xvel": - pv.get_vx(dat, out_mom=out) - elif v == "yvel": - pv.get_vy(dat, out_mom=out) - elif v == "zvel": - pv.get_vz(dat, out_mom=out) - elif v == "vel": - pv.get_vi(dat, out_mom=out) - elif v == "pressure": - pv.get_p(dat, gas_gamma=kwargs["gas_gamma"], num_moms=5, out_mom=out) - elif v == "ke": - pv.get_ke(dat, gas_gamma=kwargs["gas_gamma"], num_moms=5, out_mom=out) - elif v == "temp": - pv.get_temp(dat, gas_gamma=kwargs["gas_gamma"], num_moms=5, out_mom=out) - elif v == "sound": - pv.get_sound(dat, gas_gamma=kwargs["gas_gamma"], num_moms=5, out_mom=out) - elif v == "mach": - pv.get_mach(dat, gas_gamma=kwargs["gas_gamma"], num_moms=5, out_mom=out) - # end - # end - verb_print(ctx, "Finishing euler") diff --git a/src/postgkyl/commands/ev.py b/src/postgkyl/commands/ev.py deleted file mode 100644 index dcbbc9ef..00000000 --- a/src/postgkyl/commands/ev.py +++ /dev/null @@ -1,229 +0,0 @@ -import click -import numpy as np - -from postgkyl.commands import ev_cmd as cmd_base -from postgkyl.data import GData -from postgkyl.data import select as pselect -from postgkyl.utils import verb_print - - -help_str = "" -for s in cmd_base.cmds.keys(): - help_str += f" '{s:s}'," -# end - - -def _data(ctx, grid_stack, value_stack, ctx_stack, str_in, tags, only_active): - str_in_split = str_in.split("[") - if str_in[0] == "f" or str_in_split[0] in tags: - tag_nm = None - if str_in_split[0] in tags: - tag_nm = str_in_split[0] - only_active = False - # end - set_idx = None - if len(str_in_split) >= 2: - set_idx = str_in_split[1].split("]")[0] - # end - comp_idx = None - if len(str_in_split) == 3: - comp_idx = str_in_split[2].split("]")[0] - # end - ctx_key = None - if len(str_in.split(".")) == 2: - ctx_key = str_in.split(".")[1] - # end - - grid_stack.append([]) - value_stack.append([]) - ctx_stack.append([]) - - for dat in ctx.obj["data"].iterator(tag=tag_nm, select=set_idx, only_active=only_active): - tag_nm = dat.get_tag() - if ctx_key: - grid = None - if ctx_key in dat.ctx: - values = np.array(dat.ctx[ctx_key]) - else: - ctx.fail(click.style(f"Wrong ctx key '{ctx_key:s}' specified", fg="red")) - # end - else: - grid, values = pselect(dat, comp=comp_idx) - # end - grid_stack[-1].append(grid) - value_stack[-1].append(values) - ctx_stack[-1].append(dat.ctx) - # end - return True, (tag_nm, set_idx) - elif "(" in str_in or "[" in str_in: - value_stack.append([eval(str_in)]) - grid_stack.append([None]) - ctx_stack.append([{}]) - return True, () - elif ":" in str_in or "," in str_in: - value_stack.append([str(str_in)]) - grid_stack.append([None]) - ctx_stack.append([{}]) - return True, () - else: - try: - value_stack.append([np.array(float(str_in))]) - grid_stack.append([None]) - ctx_stack.append([{}]) - return True, () - except Exception: - return False, () - # end - # end - - -def _compare(a, b) -> bool: - if isinstance(a, np.ndarray): - return np.array_equal(a, b) - else: - return a == b - # end - - -def _command(ctx, grid_stack, value_stack, ctx_stack, str_in): - if str_in in cmd_base.cmds: - num_in = cmd_base.cmds[str_in]["num_in"] - num_out = cmd_base.cmds[str_in]["num_out"] - func = cmd_base.cmds[str_in]["func"] - else: - return False - # end - - in_grid, in_values, in_ctx, num_sets = [], [], [], [] - for i in range(num_in): - in_grid.append(grid_stack.pop()) - in_values.append(value_stack.pop()) - in_ctx.append(ctx_stack.pop()) - num_sets.append(len(in_values[-1])) - # end - for i in range(num_out): - grid_stack.append([]) - value_stack.append([]) - ctx_stack.append([]) - # end - - for set_idx in range(max(num_sets)): - tmp_grid, tmp_values, tmp_ctx = [], [], [] - for i in range(num_in): - tmp_grid.append(in_grid[i][min(set_idx, num_sets[i] - 1)]) - tmp_values.append(in_values[i][min(set_idx, num_sets[i] - 1)]) - tmp_ctx.append(in_ctx[i][min(set_idx, num_sets[i] - 1)]) - # end - try: - out_grid, out_values = func(tmp_grid, tmp_values) - except Exception as err: - ctx.fail(click.style(f"{err}", fg="red")) - # end - - # Compare the ctx data of all the inputs and copy them to a - # ctx data dictionary of the output - out_ctx = {} - remove_list = [] - for i in range(num_in): - for key in tmp_ctx[i]: - if key in out_ctx and _compare(tmp_ctx[i][key], out_ctx[key]): # tmp_ctx[i][k] == out_ctx[k]: - pass # This key has been already copied and - # matches the output; no action needed - elif key in out_ctx: - remove_list.append(key) # There is a discrepancy between - # the ctxdata; set it to remove later - else: - out_ctx[key] = tmp_ctx[i][key] # Copy the ctx data - # end - # end - # end - # Remove duplicates - remove_list = list(dict.fromkeys(remove_list)) - # Remove the discrepancies - for k in remove_list: - out_ctx.pop(k) - # end - - for i in range(num_out): - grid_stack[-num_out + i].append(out_grid[i]) - value_stack[-num_out + i].append(out_values[i]) - ctx_stack[-num_out + i].append(out_ctx) - # end - # end - return True - - -@click.command( - help=f"Manipulate datasets using math expressions. Expressions are specified using Reverse Polish Notation (RPN).\n Supported operators are: {help_str[:-1]}" -) -@click.argument("chain", nargs=1, type=click.STRING) -@click.option("--tag", "-t", help="Tag for the result") -@click.option("--label", "-l", show_default=True, help="Custom label for the result") -@click.option("--all", "-a", is_flag=True, help="Ignore the status of a dataset") -@click.pass_context -def ev(ctx, **kwargs): - verb_print(ctx, "Starting evaluate") - data = ctx.obj["data"] - - grid_stack, value_stack, ctx_stack = [], [], [] - chain_split = kwargs["chain"].split(" ") - chain_split = list(filter(None, chain_split)) - - only_active = True - if kwargs["all"]: - only_active = False - # end - - tags = list(data.tag_iterator(only_active=only_active)) - label = kwargs["label"] - if label is None: - label = kwargs["chain"] - # end - - num_datasets_in_chain = 0 - out_data_id = () - for s in chain_split: - is_data, data_id = _data(ctx, grid_stack, value_stack, ctx_stack, s, tags, only_active) - if is_data and len(data_id) > 0 and data_id != out_data_id: - num_datasets_in_chain += 1 - out_data_id = data_id - # end - if not is_data: - is_command = _command(ctx, grid_stack, value_stack, ctx_stack, s) - # end - if not is_data and not is_command: - ctx.fail(click.style(f"Evaluate input '{s:s}' represents neither data nor commad", - fg="red")) - # end - # end - - if len(value_stack) == 0: - ctx.fail(click.style("Evaluate stack is empty, there is nothing to return", fg="red")) - elif len(value_stack) > 1: - click.echo( - click.style("WARNING: Length of the evaluate stack is bigger than 1, there is a posibility of unintended behavior", - fg="yellow" )) - # end - if num_datasets_in_chain == 1 and kwargs["tag"] is None: - cnt = 0 - tag = out_data_id[0] - for out in ctx.obj["data"].iterator(tag=tag, select=out_data_id[1], only_active=only_active): - out.push(grid_stack[-1][cnt], value_stack[-1][cnt]) - cnt += 1 - # end - else: - tag = out_data_id[0] - if kwargs["tag"]: - tag = kwargs["tag"] - else: - data.deactivate_all() - # end - for grid, values, data_ctx in zip(grid_stack[-1], value_stack[-1], ctx_stack[-1]): - out = GData(tag=tag, # comp_grid=ctx.obj['compgrid'], - label=label, ctx=data_ctx) - out.push(grid, values) - data.add(out) - # end - # end - - verb_print(ctx, "Finishing ev") diff --git a/src/postgkyl/commands/ev_cmd.py b/src/postgkyl/commands/ev_cmd.py deleted file mode 100644 index 96b53087..00000000 --- a/src/postgkyl/commands/ev_cmd.py +++ /dev/null @@ -1,441 +0,0 @@ -import click -import numpy as np -from postgkyl.data.idx_parser import idx_parser - - -def _get_grid(grid0, grid1): - if grid0 is not None and grid1 is not None: - if len(grid0) > len(grid1): - return grid0 - else: - return grid1 - # end - elif grid0 is not None: - return grid0 - elif grid1 is not None: - return grid1 - else: - return None - # end - - -def add(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - out_values = in_values[0] + in_values[1] - return [out_grid], [out_values] - - -def subtract(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - out_values = in_values[1] - in_values[0] - return [out_grid], [out_values] - - -def mult(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - a, b = in_values[1], in_values[0] - if np.array_equal(a.shape, b.shape) or len(a.shape) == 0 or len(b.shape) == 0: - out_values = a * b - else: - # When multiplying phase-space and conf-space field, the - # dimensions do not match. NumPy can do a lot of things with - # broadcasting - # (https://numpy.org/doc/stable/user/basics.broadcasting.html) but - # it requires the trailing indices to match, which is opposite to - # what we have (the first indices are matching). Therefore, one can - # transpose, multiply, and transpose back... I think -- Petr Cagas - out_values = (a.transpose() * b.transpose()).transpose() - # end - return [out_grid], [out_values] - - -def dot(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - out_values = np.sum(in_values[1] * in_values[0], axis=-1)[..., np.newaxis] - return [out_grid], [out_values] - - -def divide(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - a, b = in_values[1], in_values[0] - if np.array_equal(a.shape, b.shape) or len(a.shape) == 0 or len(b.shape) == 0: - out_values = a/b - else: - # See the 'mult' comment above - out_values = (a.transpose()/b.transpose()).transpose() - # end - return [out_grid], [out_values] - - -def sqrt(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.sqrt(in_values[0]) - return [out_grid], [out_values] - - -def psin(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.sin(in_values[0]) - return [out_grid], [out_values] - - -def pcos(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.cos(in_values[0]) - return [out_grid], [out_values] - - -def ptan(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.tan(in_values[0]) - return [out_grid], [out_values] - - -def absolute(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.abs(in_values[0]) - return [out_grid], [out_values] - - -def log(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.log(in_values[0]) - return [out_grid], [out_values] - - -def log10(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.log10(in_values[0]) - return [out_grid], [out_values] - - -def minimum(in_grid, in_values): - out_values = np.atleast_1d(np.nanmin(in_values[0])) - return [[]], [out_values] - - -def minimum2(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - out_values = np.fmin(in_values[0], in_values[1]) - return [out_grid], [out_values] - - -def maximum(in_grid, in_values): - out_values = np.atleast_1d(np.nanmax(in_values[0])) - return [[]], [out_values] - - -def maximum2(in_grid, in_values): - out_grid = _get_grid(in_grid[0], in_grid[1]) - out_values = np.fmax(in_values[0], in_values[1]) - return [out_grid], [out_values] - - -def mean(in_grid, in_values): - out_values = np.atleast_1d(np.mean(in_values[0])) - return [[]], [out_values] - - -def power(in_grid, in_values): - out_grid = in_grid[1] - out_values = np.power(in_values[1], in_values[0]) - return [out_grid], [out_values] - - -def sq(in_grid, in_values): - out_grid = in_grid[0] - out_values = in_values[0]**2 - return [out_grid], [out_values] - - -def exp(in_grid, in_values): - out_grid = in_grid[0] - out_values = np.exp(in_values[0]) - return [out_grid], [out_values] - - -def length(in_grid, in_values): - ax = int(in_values[0]) - ln = in_grid[1][ax][-1] - in_grid[1][ax][0] - if len(in_grid[1][ax]) == in_values[1].shape[ax]: - ln += in_grid[1][ax][1] - in_grid[1][ax][0] - # end - return [[]], [ln] - - -def grad(in_grid, in_values): - out_grid = in_grid[0] - nd = len(in_values[0].shape) - 1 - out_shape = list(in_values[0].shape) - nc = in_values[0].shape[-1] - out_shape[-1] = nc * nd - out_values = np.zeros(out_shape) - - for d in range(nd): - zc = 0.5 * (in_grid[0][d][1:] + in_grid[0][d][:-1]) # get cell centered values - out_values[..., d*nc:(d + 1)*nc] = np.gradient( - in_values[0], zc, edge_order=2, axis=d - ) - # end - return [out_grid], [out_values] - - -def grad2(in_grid, in_values): - out_grid = in_grid[1] - ax = in_values[0] - if isinstance(ax, str) and ":" in ax: - tmp = ax.split(":") - lo = int(tmp[0]) - up = int(tmp[1]) - rng = range(lo, up) - elif isinstance(ax, str): - rng = tuple((int(i) for i in ax.split(","))) - else: - rng = range(int(ax), int(ax + 1)) - # end - - num_dims = len(rng) - out_shape = list(in_values[1].shape) - num_comps = in_values[1].shape[-1] - out_shape[-1] = out_shape[-1] * num_dims - out_values = np.zeros(out_shape) - - for cnt, d in enumerate(rng): - zc = 0.5 * (in_grid[1][d][1:] + in_grid[1][d][:-1]) # get cell centered values - out_values[..., cnt*num_comps:(cnt + 1)*num_comps] = np.gradient( - in_values[1], zc, edge_order=2, axis=d - ) - # end - return [out_grid], [out_values] - - -def integrate(in_grid, in_values, avg=False): - grid = in_grid[1].copy() - values = np.array(in_values[1]) - - axis = in_values[0] - if isinstance(axis, float): - axis = tuple([int(axis)]) - elif isinstance(axis, tuple): - pass - elif isinstance(axis, np.ndarray): - axis = tuple([int(axis)]) - elif isinstance(axis, str): - if len(axis.split(",")) > 1: - axes = axis.split(",") - axis = tuple([int(a) for a in axes]) - elif len(axis.split(":")) == 2: - bounds = axis.split(":") - axis = tuple(range(bounds[0], bounds[1])) - elif axis == "all": - num_dims = len(grid) - axis = tuple(range(num_dims)) - # end - else: - raise TypeError("'axis' needs to be integer, tuple, string of comma separated integers, or a slice ('int:int')") - # end - - dz = [] - for d, coord in enumerate(grid): - dz.append(coord[1:] - coord[:-1]) - if len(coord) == values.shape[d]: - dz[-1] = np.append(dz[-1], dz[-1][-1]) - # end - - # Integration assuming values are cell centered averages - # Should work for nonuniform meshes - for ax in sorted(axis, reverse=True): - values = np.moveaxis(values, ax, -1) - values = np.dot(values, dz[ax]) - # end - for ax in sorted(axis): - grid[ax] = np.array([0]) - values = np.expand_dims(values, ax) - if avg: - ln = in_grid[1][ax][-1] - in_grid[1][ax][0] - if len(in_grid[1][ax]) == in_values[1].shape[ax]: - ln += in_grid[1][ax][1] - in_grid[1][ax][0] - # end - values = values/ln - # end - # end - return [grid], [values] - - -def average(in_grid, in_values): - return integrate(in_grid, in_values, True) - - -def divergence(in_grid, in_values): - out_grid = in_grid[0] - num_dims = len(in_grid[0]) - num_comps = in_values[0].shape[-1] - if num_comps > num_dims: - click.echo( - click.style(f"WARNING in 'ev div': Length of the provided vector ({num_comps:d}) is longer than number of dimensions ({num_dims:d}). The last {num_comps - num_dims:d} component(s) of the vector will be disregarded.", - fg="yellow") - ) - # end - out_shape = list(in_values[0].shape) - out_shape[-1] = 1 - out_values = np.zeros(out_shape) - for d in range(num_dims): - zc = 0.5 * (in_grid[0][d][1:] + in_grid[0][d][:-1]) # get cell centered values - out_values[..., 0] = out_values[..., 0] + np.gradient( - in_values[0][..., d], zc, edge_order=2, axis=d - ) - # end - return [out_grid], [out_values] - - -def curl(in_grid, in_values): - out_grid = in_grid[0] - num_dims = len(in_grid[0]) - num_comps = in_values[0].shape[-1] - - out_shape = list(in_values[0].shape) - - if num_dims == 1: - if num_comps != 3: - raise ValueError(f"ERROR in 'ev curl': Curl in 1D requires 3-component input and {num_comps:d}-component field was provided.") - # end - zc0 = 0.5*(in_grid[0][0][1:] + in_grid[0][0][:-1]) - out_values = np.zeros(out_shape) - out_values[..., 1] = -np.gradient(in_values[0][..., 2], zc0, edge_order=2, axis=0) - out_values[..., 2] = np.gradient(in_values[0][..., 1], zc0, edge_order=2, axis=0) - elif num_dims == 2: - zc0 = 0.5 * (in_grid[0][0][1:] + in_grid[0][0][:-1]) - zc1 = 0.5 * (in_grid[0][1][1:] + in_grid[0][1][:-1]) - if num_comps < 2: - raise ValueError(f"ERROR in 'ev curl': Length of the provided vector ({num_comps:d}) is smaller than number of dimensions ({num_dims:d}). Curl can't be calculated." ) - elif num_comps == 2: - click.echo( - click.style(f"WARNING in 'ev curl': Length of the provided vector ({num_comps:d}) is longer than number of dimensions ({num_dims:d}). Only the third component of curl will be calculated.", - fg="yellow") - ) - out_shape[-1] = 1 - out_values = np.zeros(out_shape) - out_values[..., 0] = np.gradient( - in_values[0][..., 1], zc0, edge_order=2, axis=0 - ) - np.gradient(in_values[0][..., 0], zc1, edge_order=2, axis=1) - else: - if num_comps > 3: - print("here") - click.echo( - click.style(f"WARNING in 'ev curl': Length of the provided vector ({num_comps:d}) is longer than number of dimensions ({num_dims:d}). The last {num_comps - num_dims:d} components of the vector will be disregarded.", - fg="yellow") - ) - # end - out_values = np.zeros(out_shape) - out_values[..., 0] = np.gradient(in_values[0][..., 2], zc1, edge_order=2, axis=1) - out_values[..., 1] = -np.gradient(in_values[0][..., 2], zc0, edge_order=2, axis=0) - out_values[..., 2] = np.gradient( in_values[0][..., 1], zc0, edge_order=2, axis=0) - np.gradient(in_values[0][..., 0], zc1, edge_order=2, axis=1) - else: # 3D - if num_comps > 3: - click.echo( - click.style(f"WARNING in 'ev curl': Length of the provided vector ({num_comps:d}) is longer than number of dimensions ({num_dims:d}). The last {num_comps - num_dims:d} component(s) of the vector will be disregarded.", - fg="yellow") - ) - elif num_comps < 3: - raise ValueError( - f"ERROR in 'ev curl': Length of the provided vector ({num_comps:d}) is smaller than number of dimensions ({num_dims:d}). Curl can't be calculated." - ) - # end - zc0 = 0.5 * (in_grid[0][0][1:] + in_grid[0][0][:-1]) - zc1 = 0.5 * (in_grid[0][1][1:] + in_grid[0][1][:-1]) - zc2 = 0.5 * (in_grid[0][2][1:] + in_grid[0][2][:-1]) - out_values[..., 0] = np.gradient(in_values[0][..., 2], zc1, edge_order=2, axis=1) - np.gradient(in_values[0][..., 1], zc2, edge_order=2, axis=2) - out_values[..., 1] = np.gradient(in_values[0][..., 0], zc2, edge_order=2, axis=2) - np.gradient(in_values[0][..., 2], zc0, edge_order=2, axis=0) - out_values[..., 2] = np.gradient(in_values[0][..., 1], zc0, edge_order=2, axis=0) - np.gradient(in_values[0][..., 0], zc1, edge_order=2, axis=1) - # end - return [out_grid], [out_values] - -def scale_comp(in_grid, in_values): - """Scale specific components of the data. - - Args: - in_values[0]: Scaling factor (float) - from RPN stack order - in_values[1]: Component specification (string like "2:4" or number) - in_values[2]: Original data array (f) - - Usage: f 2:4 1000 scale_comp (scales components 2 and 3 by 1000) - """ - - out_grid = in_grid[2] # Use grid from original data (f) - original_data = in_values[2].copy() # Original data (make a copy) - comp_spec = in_values[1] # Component specification (can be string or number) - scale_factor = in_values[0] # Scaling factor - - scale_factor = scale_factor.item() # Ensure scale_factor is a float - # Parse component specification - if isinstance(comp_spec, str): - comp_idx = idx_parser(comp_spec) - elif isinstance(comp_spec, np.ndarray) and comp_spec.size == 1: - # Handle single number case - comp_idx = int(comp_spec.item()) - else: - comp_idx = int(comp_spec) - - # Apply scaling to specified components - if isinstance(comp_idx, slice): - original_data[..., comp_idx] *= scale_factor - elif isinstance(comp_idx, tuple): - for idx in comp_idx: - original_data[..., idx] *= scale_factor - else: - original_data[..., comp_idx] *= scale_factor - - return [out_grid], [original_data] - -def scale_zi_axis(in_grid, in_values): - """Scale the axis of the z_i dimension of the data - - Args: - in_values[0]: Scaling factor (float) - from RPN stack order - in_values[1]: Axis direction (int) - 0,1,2,3,4,5 - in_values[2]: Original data array (f) - - Usage: f 1000 scale_xaxis (scales x-axis by 1000) - """ - - out_grid = in_grid[2] # Use grid from original data (f) - original_data = in_values[2].copy() # Original data (make a copy) - idx_scale = in_values[1].item() # Axis direction (int) - scale_factor = in_values[0].item() # Ensure scale_factor is a float - - # Scale the z_i axis - out_grid[int(idx_scale)] *= scale_factor - - return [out_grid], [original_data] - -cmds = { - "+": {"num_in": 2, "num_out": 1, "func": add}, - "-": {"num_in": 2, "num_out": 1, "func": subtract}, - "*": {"num_in": 2, "num_out": 1, "func": mult}, - "/": {"num_in": 2, "num_out": 1, "func": divide}, - "dot": {"num_in": 2, "num_out": 1, "func": dot}, - "sqrt": {"num_in": 1, "num_out": 1, "func": sqrt}, - "sin": {"num_in": 1, "num_out": 1, "func": psin}, - "cos": {"num_in": 1, "num_out": 1, "func": pcos}, - "tan": {"num_in": 1, "num_out": 1, "func": ptan}, - "abs": {"num_in": 1, "num_out": 1, "func": absolute}, - "avg": {"num_in": 2, "num_out": 1, "func": average}, - "log": {"num_in": 1, "num_out": 1, "func": log}, - "log10": {"num_in": 1, "num_out": 1, "func": log10}, - "max": {"num_in": 1, "num_out": 1, "func": maximum}, - "min": {"num_in": 1, "num_out": 1, "func": minimum}, - "max2": {"num_in": 2, "num_out": 1, "func": maximum2}, - "min2": {"num_in": 2, "num_out": 1, "func": minimum2}, - "mean": {"num_in": 1, "num_out": 1, "func": mean}, - "len": {"num_in": 2, "num_out": 1, "func": length}, - "pow": {"num_in": 2, "num_out": 1, "func": power}, - "sq": {"num_in": 1, "num_out": 1, "func": sq}, - "exp": {"num_in": 1, "num_out": 1, "func": exp}, - "grad": {"num_in": 1, "num_out": 1, "func": grad}, - "grad2": {"num_in": 2, "num_out": 1, "func": grad2}, - "int": {"num_in": 2, "num_out": 1, "func": integrate}, - "div": {"num_in": 1, "num_out": 1, "func": divergence}, - "curl": {"num_in": 1, "num_out": 1, "func": curl}, - "scale_comp": {"num_in": 3, "num_out": 1, "func": scale_comp}, - "scale_zi_axis": {"num_in": 3, "num_out": 1, "func": scale_zi_axis}, -} diff --git a/src/postgkyl/commands/extractinput.py b/src/postgkyl/commands/extractinput.py deleted file mode 100644 index ec7719a1..00000000 --- a/src/postgkyl/commands/extractinput.py +++ /dev/null @@ -1,24 +0,0 @@ -import base64 -import click - -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.pass_context -def extractinput(ctx, **kwargs): - """Extract embedded input file from compatible BP files""" - verb_print(ctx, "Starting ") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - enc_inp = dat.get_input_file() - if enc_inp: - inpfile = base64.decodebytes(enc_inp.encode("utf-8")).decode("utf-8") - click.echo(inpfile) - else: - click.echo("No embedded input file!") - # end - # end - verb_print(ctx, "Finishing extractinput") diff --git a/src/postgkyl/commands/fft.py b/src/postgkyl/commands/fft.py deleted file mode 100644 index 020d3f5c..00000000 --- a/src/postgkyl/commands/fft.py +++ /dev/null @@ -1,37 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.fft - - -@click.command() -@click.option("-p", "--psd", is_flag=True, - help="Limits output to positive frequencies and returns the power spectral density |FT|^2.") -@click.option("-i", "--iso", is_flag=True, - help="Bins power spectral density |FT|^2, making 1D power spectra from multi-dimensional data.") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.pass_context -def fft(ctx, **kwargs): - """Calculate the Fourier Transform or the power-spectral density of input data. - - Only works on 1D data at present. - """ - verb_print(ctx, "Starting FFT") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - grid, values = postgkyl.tools.fft(dat, psd=kwargs["psd"], iso=kwargs["iso"]) - out.push(grid, values) - data.add(out) - else: - postgkyl.tools.fft(dat, psd=kwargs["psd"], iso=kwargs["iso"], overwrite=True) - # end - # end - - verb_print(ctx, "Finishing FFT") diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py deleted file mode 100644 index 77d40361..00000000 --- a/src/postgkyl/commands/gk_distf.py +++ /dev/null @@ -1,244 +0,0 @@ -import glob - -import click -import numpy as np - -from postgkyl.data import GData, GInterpModal -from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops -from postgkyl.utils import verb_print - -# mc2nu grid deformation helpers -# This is a result of the gkyl_reader not having support for both mapc2p and mapc2p-vel grids. -# Particularly, the gkyl_reader does not support mapping phase space arrays with mapc2p -def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: - """ - Given an array defined at cell centers, return the corresponding nodal - values by interpolating half a cell width at the boundaries. - """ - nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) - nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) - nodes[0] = cell_centers[0] + (cell_centers[0] - nodes[1]) # Cell center plus half a cell width - nodes[-1] = cell_centers[-1] + (cell_centers[-1] - nodes[-2]) # Cell center plus half a cell width - return nodes - -def _extract_values_along_dimension(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: - """Decompose mapped_values into a 1D array along the specified axis""" - idx = [0] * (cdim + 1) # Initialize indexing array. mc2nu has cdim+1 dimensions. - idx[axis] = slice(None) # Define a slice along the desired axis. - idx[-1] = axis # Select the appropriate component of mc2nu - return mapped_values[tuple(idx)].reshape(-1) # Apply indices and flatten to 1D. - -def _apply_mc2nu_grid(uniform_grid: list, mc2nu_file: str, interp: int | None = None) -> list: - """Replace computational configuration-space grid with non-uniform spatial coordinates.""" - mc2nu_data = GData(mc2nu_file) - cdim = mc2nu_data.get_num_dims() - - _, mc2nu_values = GInterpModal(mc2nu_data, 1, "ms", interp).interpolate(tuple(range(cdim))) - - nonuniform_grid = list(uniform_grid) - for d in range(cdim): - mc2nu_single_axis = _extract_values_along_dimension(mc2nu_values, d, cdim) - nonuniform_grid[d] = _convert_cell_centered_to_nodal(mc2nu_single_axis) - # end - return nonuniform_grid - -def _resolve_optional_file_option(option_value: str | None) -> tuple[bool, str | None]: - """Interpret an optional-value CLI option as (enabled, override_file).""" - if option_value is None: - return False, None - if option_value == "": - return True, None - return True, option_value - -def load_gk_distf( - name: str, species: str, frame: int, - tag: str = "f", suffix: str = "", use_c2p_vel: bool = False, - use_mc2nu: bool = False, use_mapc2p: bool = False, block_idx: int | None = None, - interp: int | None = None, - Jf_file: str | None = None, - mapc2p_vel_file: str | None = None, - jacobvel_file: str | None = None, - mc2nu_file: str | None = None, - mapc2p_file: str | None = None, - jacobtot_inv_file: str | None = None, ) -> GData: - """Build a real distribution function from saved JxJvBf data.""" - - prefix = f"{name}_b{block_idx}" if block_idx is not None else name - frame_infix = f"{suffix}_" if suffix else "" - - if Jf_file is None: - Jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" - # end - if mapc2p_vel_file is None: - mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" - # end - if jacobvel_file is None: - jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" - # end - if mc2nu_file is None: - mc2nu_file = f"{prefix}-geo_corn_mc2nu_pos_deflated.gkyl" - # end - if mapc2p_file is None: - mapc2p_file = f"{prefix}-geo_corn_mapc2p_deflated.gkyl" - # end - if jacobtot_inv_file is None: - jacobtot_inv_file = f"{prefix}-geo_int_jacobtot_inv.gkyl" - # end - - Jf_data = GData(Jf_file, mapc2p_vel_name=mapc2p_vel_file if use_c2p_vel else None) - jacobvel_data = GData(jacobvel_file) - jacobtot_inv_data = GData(jacobtot_inv_file) - - # Divide Jf by jacobvel to get f * J_x * B. - fJxB_data = GData(ctx=Jf_data.ctx) # Inside a GData object so we can interpolate - fJxB_values = Jf_data.get_values() / jacobvel_data.get_values() - fJxB_data.push(Jf_data.get_grid(), fJxB_values) - - if interp == 0: - # No interpolation: weak multiply by reciprocal of (J_x*B). - out_grid = fJxB_data.get_grid() - f_data = GData(ctx=fJxB_data.ctx) - f_data.push(out_grid, np.zeros_like(fJxB_data.get_values())) - GkeyllDGops().multiply_conf_phase(f_data, jacobtot_inv_data, fJxB_data) - f_values = f_data.get_values() - else: - # Interpolate f * J_x * B and jacobtot_inv to the same grid. - out_grid, fJxB_values = GInterpModal(fJxB_data, 1, "gkhyb", interp).interpolate() - _, jacobtot_inv_values = GInterpModal(jacobtot_inv_data, 1, "ms", interp).interpolate() - fJxB_values = np.squeeze(fJxB_values) - jacobtot_inv_values = np.squeeze(jacobtot_inv_values) - - # Reshape jacobtot_inv to have 1 component over velocity dimensions, then multiply. - vdim = fJxB_values.ndim - jacobtot_inv_values.ndim - jacobtot_inv_reshaped = jacobtot_inv_values.reshape(jacobtot_inv_values.shape + (1,) * vdim) - f_values = fJxB_values * jacobtot_inv_reshaped - # Add 1 dimension to represent 1 component - f_values = f_values.reshape(f_values.shape + (1,)) - - if use_mc2nu: - out_grid = _apply_mc2nu_grid(out_grid, mc2nu_file, interp) - if use_c2p_vel: - Jf_data.ctx["grid_type"] = "c2p_vel + mc2nu" - else: - Jf_data.ctx["grid_type"] = "mc2nu" - # end - elif use_mapc2p: - out_grid = _apply_mc2nu_grid(out_grid, mapc2p_file, interp) - if use_c2p_vel: - Jf_data.ctx["grid_type"] = "c2p_vel + mapc2p" - else: - Jf_data.ctx["grid_type"] = "mapc2p" - # end - # end - - out = GData(tag=tag, ctx=Jf_data.ctx) - out.push(out_grid, f_values) - return out -# end - -@click.command() -@click.option("--name", "-n", required=True, type=click.STRING, - help="Simulation name prefix (e.g. gk_lorentzian_mirror).") -@click.option("--species", "-s", required=True, type=click.STRING, - help="Species name (e.g. ion or elc).") -@click.option("--suffix", default="", type=click.STRING, - help="Use -__.gkyl as the input distribution.") -@click.option("--Jf-file", default=None, type=click.STRING, - help="Jf filename override. If omitted, the default naming convention is used.") -@click.option("--jacobvel-file", default=None, type=click.STRING, - help="jacobvel filename override. If omitted, the default naming convention is used.") -@click.option("--jacobtot-inv-file", default=None, type=click.STRING, - help="jacobtot_inv filename override. If omitted, the default naming convention is used.") -@click.option("--frame", "-f", required=True, type=click.STRING, - help="Frame number, comma separated values, or range. Use ':' for all frames\n" - " and 'start:stop[:step]' for ranges.") -@click.option("--interp", "-i", type=click.INT, - help="Interpolation onto a general mesh of specified amount. User -i 0 for no interpolation.") -@click.option("--c2p-vel", "-v", default=None, flag_value="", type=click.STRING, - help="Convert velocity-space computational to physical coordinates, using mapping\n" - "in (optionally) given file (default *_mapc2p_vel.gkyl).") -@click.option("--mc2nu", "-m", default=None, flag_value="", type=click.STRING, - help="Convert non-uniform computational to field-aligned coordinates using mapping \n" - "in (optionally) given file (default: *-geo_corn_mc2nu_pos_deflated.gkyl).") -@click.option("--mapc2p", "-p", default=None, flag_value="", type=click.STRING, - help="Convert position-space computational to Cartesian (GKYL_GEOMETRY_MAPC2P) or \n" - "cylindrical (GKYL_GEOMETRY_TOKAMAK, GKYL_GEOMETRY_MIRROR) coordinates, using \n" - "mapping in (optionally) given file (default: *-geo_corn_mapc2p.gkyl)") -@click.option("--block", "-b", default=None, type=click.INT, - help="Use block-specific files with _b prefix, e.g. -b 1 loads _b1-*.gkyl.") -@click.option("--tag", "-t", default="f", type=click.STRING, - help="Tag for output dataset.") -@click.pass_context -def gk_distf(ctx, **kwargs): - """ - Gyrokinetics: load the distribution function from files containing the - distribution (f) times one or multiple Jacobians (J). The Jacobians are - divided out in order to output f. The distribution is interpolated, and - the interpolation can optionally use mappings to convert from computational - to physical coordinates. - - \b - Command line example: - pgkyl gk-distf -n gk_lorentzian_mirror -s ion -f 0 - - \b - Script example: - import postgkyl as pg - from postgkyl.commands import load_gk_distf - - distf = pg.commands.load_gk_distf(name="gk_lorentzian_mirror", species="ion", frame=0) - """ - data = ctx.obj["data"] - - verb_print(ctx, "Building distribution function for " + kwargs["name"]) - - frame_spec = kwargs["frame"].strip() - if "," in frame_spec: - frames = [int(f.strip()) for f in frame_spec.split(",")] # List of frames specified on input - elif ":" not in frame_spec: - frames = [int(frame_spec)] # Stick to the frame specified on input - else: - # Figure out how many frames are possible to read based on what files are available - prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] - frame_infix = f"{kwargs['suffix']}_" if kwargs["suffix"] else "" - stem = f"{prefix}-{kwargs['species']}_{frame_infix}" - available = sorted({ - int(f.removeprefix(stem)[:-5]) - for f in glob.glob(f"{glob.escape(stem)}*.gkyl") - if f.removeprefix(stem)[:-5].isdigit() - }) - # Slice the data accordingly - parts = frame_spec.split(":") - lower = int(parts[0]) if parts[0] else available[0] - upper = int(parts[1]) if parts[1] else available[-1] + 1 - step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 - frames = [f for f in available if lower <= f < upper and (f - lower) % step == 0] - # end - verb_print(ctx, f"Loading frames: {frames}") - - use_c2p_vel, mapc2p_vel_file = _resolve_optional_file_option(kwargs["c2p_vel"]) - use_mc2nu, mc2nu_file = _resolve_optional_file_option(kwargs["mc2nu"]) - use_mapc2p, mapc2p_file = _resolve_optional_file_option(kwargs["mapc2p"]) - - for frame in frames: - out = load_gk_distf( - name=kwargs["name"], species=kwargs["species"], frame=frame, - tag=kwargs["tag"], suffix=kwargs["suffix"], - use_c2p_vel=use_c2p_vel, - use_mc2nu=use_mc2nu, use_mapc2p=use_mapc2p, - block_idx=kwargs["block"], - interp=kwargs["interp"], - Jf_file=kwargs.get("Jf-file"), - mapc2p_vel_file=mapc2p_vel_file, - jacobvel_file=kwargs["jacobvel_file"], - mc2nu_file=mc2nu_file, - mapc2p_file=mapc2p_file, - jacobtot_inv_file=kwargs["jacobtot_inv_file"], - ) - data.add(out) - # end - - if len(frames) > 1: - data.set_unique_labels() - # end -# end diff --git a/src/postgkyl/commands/gk_energy_balance.py b/src/postgkyl/commands/gk_energy_balance.py deleted file mode 100644 index b3d8f26c..00000000 --- a/src/postgkyl/commands/gk_energy_balance.py +++ /dev/null @@ -1,527 +0,0 @@ -import click -import numpy as np -import matplotlib.pyplot as plt -import os -import glob - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -@click.command() -@click.option("--name", "-n", required=True, type=click.STRING, default=None, - help="Simulation name (also the file prefix, e.g. gk_sheath_1x2v_p1).") -@click.option("--species", "-s", required=True, default=None, - help="Comma-separated list of species names.") -@click.option("--path", "-p", type=click.STRING, default='./', - help="Path to simulation data.") -@click.option("--relative_error", "-r", is_flag=True, - help="Plot the relative error only.") -@click.option("--multib", "-m", is_flag=False, default="-10", flag_value="-1", - help="Multiblock. Optional: pass block indices as comma-separated list or slice (start:stop:step). If no indices are given, all blocks are used.") -@click.option("--field_dot_file", type=click.STRING, default=None, multiple=True, - help="Integrated field energy rate of change.") -@click.option("--apar_dot_file", type=click.STRING, default=None, multiple=True, - help="Integrated apar energy rate of change.") -@click.option("--fdot_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of change in f over a time step.") -@click.option("--source_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of the source(s).") -@click.option("--bflux_xlower_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through lower x boundary.") -@click.option("--bflux_ylower_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through lower y boundary.") -@click.option("--bflux_zlower_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through lower z boundary.") -@click.option("--bflux_xupper_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through upper x boundary.") -@click.option("--bflux_yupper_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through upper y boundary.") -@click.option("--bflux_zupper_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through upper z boundary.") -@click.option("--f_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of f.") -@click.option("--field_file", type=click.STRING, default=None, multiple=True, - help="Integrated field energy.") -@click.option("--apar_file", type=click.STRING, default=None, multiple=True, - help="Integrated apar energy.") -@click.option("--dt_file", type=click.STRING, default=None, - help="Time step.") -@click.option("--logy", is_flag=True, default=False, - help="Logarithmic scale for y axis.") -@click.option("--absy", is_flag=True, default=False, - help="Take absolute value of time traces.") -@click.option("--xlabel", type=click.STRING, default="Time (s)", - help="Label for the x axis.") -@click.option("--ylabel", type=click.STRING, default=None, - help="Label for the y axis.") -@click.option("--title", type=click.STRING, default=None, - help="Take absolute value of time traces.") -@click.option("--indent_left", type=click.FLOAT, default=0.0, - help="A number in the [-0.11,0.88] range by which to shift the left boundary of the plot.") -@click.option("--add_width", type=click.FLOAT, default=0.0, - help="A number in the [-0.86,0.13] range by which to increase the width the plot.") -@click.option("--saveas", type=click.STRING, default=None, - help="Name of figure file.") -@click.pass_context -def gk_energy_balance(ctx, **kwargs): - """ - \b - Gyrokinetics: Plot the energy balance of a simulation. - Requires the following files: - -field_energy_dot.gkyl - ..._fdot_integrated_moms.gkyl - ..._source_integrated_moms.gkyl - ..._bflux__integrated_HamiltonianMoments.gkyl - where ... means -, and we need these - files for each species. The last two files above are only needed if - the simulation had sources or non-periodic boundaries. - For electromagnetic simulations, the following file is also used - (if present): - -apar_energy_dot.gkyl - If the relative error is requested, these are also needed: - ..._integrated_moms.gkyl - -field_energy.gkyl - -apar_energy.gkyl (electromagnetic only) - -dt.gkyl - - \b - The default assumes these are in the current directory. - Alternatively, the full path to each file can be specified. - If passing the full path for the species-specific filed (e.g. --fdot_file) - pass * for the species name. - - \b - If simulation is multiblock, and you wish to specify files manually: - 1) Pass * for the block index. - 2) Use --multib/-m to specify desired blocks (or ommit to use all). - - NOTE: this command cannot be combined with other postgkyl commands. - """ - - # - # Hardcoded parameters and auxiliary functions. - # - max_num_blocks = 10000 - - # Labels used to identify boundary flux files. - edges = ["lower","upper"] - dirs = ["x","y","z"] - # Line styles. - line_styles = ['-','--',':','-.','None','None','None','None'] - # Font sizes. - xy_label_font_size = 17 - title_font_size = 17 - tick_font_size = 14 - legend_font_size = 14 - - # Create figure. - figProp1a = (7.5, 4.5) - ax1aPos = [0.11+kwargs["indent_left"], 0.15, 0.87+kwargs["add_width"], 0.78] - fig1a = plt.figure(figsize=figProp1a) - ax1a = fig1a.add_axes(ax1aPos) - - def set_tick_font_size(axIn,fontSizeIn): - # Set the font size of the ticks to a given size. - axIn.tick_params(axis='both',labelsize=fontSizeIn) - offset_txt = axIn.yaxis.get_offset_text() # Get the text object - offset_txt.set_size(fontSizeIn) # Set the size. - offset_txt = axIn.xaxis.get_offset_text() # Get the text object - offset_txt.set_size(fontSizeIn) # Set the size. - - def read_gfile_if_present(file_name): - # Check if a Gkeyll file exists. If it does, read it and return - # its grid, data and GData object. If it doesn't, return None. - if os.path.exists(file_name): - pgData = GData(file_name) # Read data with pgkyl. - time = pgData.get_grid() # Time stamps of the simulation. - val = pgData.get_values() # Data values. - return True, np.squeeze(time), np.squeeze(val), pgData - else: - verb_print(ctx, " -> File "+file_name+" not found. Proceeding w/o it.") - return False, None, None, None - - def parse_slice_string(value): - # Parse a 'slice()' from string, like 'start:stop:step'. - parts = value.split(':') - # Convert parts to integers, replacing empty strings with None for slice defaults - parsed_parts = [] - for p in parts: - try: - parsed_parts.append(int(p) if p else None) - except ValueError: - # Handle cases where the part might not be a number - raise ValueError(f"Invalid slice part: {p}") - # Create the slice object with the appropriate number of arguments - return slice(*parsed_parts) - - def accumulate_or_assign(target_arr, old_arr): - # Accumulates old_arr into target_arr if target_arr exists, - # otherwise assign old_arr to target_arr. - old_arr = np.asarray(old_arr) # Ensure old_arr is a numpy array. - if target_arr is None: - return old_arr.copy() - else: - target_arr += old_arr - return target_arr - - def absy_enabled(data_in): - # Take the absolute value of the data - return np.abs(data_in) - - def absy_disabled(data_in): - # Don't take the absolute value of the data - return data_in - # - # End of hardcoded parameters and auxiliary functions. - # - - data = ctx.obj["data"] # Data stack. - - verb_print(ctx, "Plotting energy balance for " + kwargs["name"]) - - absy_func = absy_disabled - if kwargs["absy"]: - absy_func = absy_enabled - - kwargs["path"] = kwargs["path"] + '/' # For safety. - - species_names = kwargs["species"].split(",") # Name of species simulated. - num_species = len(species_names) - - # Determine blocks to plot, number of blocks, and set file prefix. - if kwargs["multib"] == "-10": - # Single block. - file_path_prefix = kwargs["path"] + kwargs["name"] + '-' - blocks = [0] - num_blocks = 1 - else: - # Multi block. - file_path_prefix = kwargs["path"] + kwargs["name"] + '_b*-' - - if kwargs["multib"] == "-1": - # Find and use all blocks. - if kwargs["fdot_file"]: - fdot_file = kwargs["path"] + kwargs["fdot_file"] - fdot_file = fdot_file[::-1].replace("*",species_name[0],1)[::-1] - else: - fdot_file = file_path_prefix + species_names[0] + '_fdot_integrated_moms.gkyl' - - fdot_file_list = glob.glob(fdot_file) - num_blocks = len(fdot_file_list) - blocks = list(range(num_blocks)) - else: - # Use specified blocks. - if ',' in kwargs["multib"]: - blocks = kwargs["multib"].split(",") - num_blocks = len(blocks) - blocks = [int(blocks[i]) for i in range(num_blocks)] - elif ':' in kwargs["multib"]: - slice_obj = parse_slice_string(kwargs["multib"]) - max_num_blocks = 10000 - blocks = list(range(*slice_obj.indices(max_num_blocks))) - num_blocks = len(blocks) - - else: - raise NameError("Blocks given to --multib -m must be a comma separated list or slice.") - - block_path_prefix = file_path_prefix - - field_dot = None - apar_dot = None - fdot = None - src = None - bflux_tot = None - for bI in range(num_blocks): - - block_path_prefix = file_path_prefix.replace("*",str(bI)) - - # Load field energy rate of change data. - if kwargs["field_dot_file"]: - field_dot_file = kwargs["path"] + kwargs["field_dot_file"].replace("*",str(bI)) - else: - field_dot_file = block_path_prefix + 'field_energy_dot.gkyl' - - has_field_dot, time_field_dot, field_dot_pb, gdat = read_gfile_if_present(field_dot_file) - if not has_field_dot or gdat is None: - raise FileNotFoundError(f"Required file not found: {field_dot_file}") - gdat_field_dot = GData(tag="field_dot", label="field_dot", ctx=gdat.ctx) - - # Load apar energy rate of change data (optional, may not exist in electrostatic simulations). - if kwargs["apar_dot_file"]: - apar_dot_file = kwargs["path"] + kwargs["apar_dot_file"].replace("*",str(bI)) - else: - apar_dot_file = block_path_prefix + 'apar_energy_dot.gkyl' - - has_apar_dot, time_apar_dot, apar_dot_pb, gdat = read_gfile_if_present(apar_dot_file) - if has_apar_dot: - gdat_apar_dot = GData(tag="apar_dot", label="apar_dot", ctx=gdat.ctx) - - fdot_pb = None - src_pb = None - bflux_tot_pb = None - for sI in range(len(species_names)): - spec_nm = species_names[sI] - - # Load change in species over a time step. - if kwargs["fdot_file"]: - fdot_file = (kwargs["path"] + kwargs["fdot_file"].replace("*",str(bI),1)).replace("*",spec_nm) - else: - fdot_file = block_path_prefix + spec_nm + '_fdot_integrated_moms.gkyl' - - has_fdot, time_fdot, fdot_ps, gdat = read_gfile_if_present(fdot_file) - if not has_fdot or gdat is None: - raise FileNotFoundError(f"Required file not found: {fdot_file}") - gdat_fdot = GData(tag="fdot", label="fdot", ctx=gdat.ctx) - - # Load integrated moments of the source. - if kwargs["source_file"]: - src_file = (kwargs["path"] + kwargs["source_file"].replace("*",str(bI),1)).replace("*",spec_nm) - else: - src_file = block_path_prefix + spec_nm + '_source_integrated_moms.gkyl' - - has_src, time_src, src_ps, gdat = read_gfile_if_present(src_file) - if has_src: - gdat_src = GData(tag="src", label="src", ctx=gdat.ctx) - - # Load particle boundary fluxes. - nbflux = 0 - time_bflux, bflux_ps = list(), list() - has_bflux = False - for d in dirs: - for e in edges: - if kwargs["bflux_"+d+e+"_file"]: - bflux_file = (kwargs["path"] + kwargs["bflux_"+d+e+"_file"].replace("*",str(bI),1)).replace("*",spec_nm) - else: - bflux_file = block_path_prefix + spec_nm + '_bflux_'+d+e+'_integrated_HamiltonianMoments.gkyl' - - has_bflux_at_boundary, time_bflux_tmp, bflux_tmp, gdat = read_gfile_if_present(bflux_file) - if has_bflux_at_boundary: - gdat_bflux = GData(tag="bflux", label="bflux", ctx=gdat.ctx) - time_bflux.append(time_bflux_tmp) - bflux_ps.append(bflux_tmp) - has_bflux = has_bflux or has_bflux_at_boundary - nbflux += 1 - - #[ Select the Hamiltonian moment. - fdot_ps = fdot_ps[:,2] - if has_src: - src_ps = src_ps[:,2] - else: - src_ps = 0.0*fdot_ps - - if has_bflux: - for i in range(nbflux): - bflux_ps[i] = bflux_ps[i][:,2] - - # Add boundary fluxes of all boundaries. - if has_bflux: - time_bflux_tot = time_bflux[0] - bflux_tot_ps = bflux_ps[0] - for i in range(1,nbflux): - bflux_tot_ps += bflux_ps[i] - else: - bflux_tot_ps = 0.0*fdot_ps - - # Add over species. - fdot_pb = accumulate_or_assign(fdot_pb, fdot_ps) - src_pb = accumulate_or_assign(src_pb, src_ps) - bflux_tot_pb = accumulate_or_assign(bflux_tot_pb, bflux_tot_ps) - - # Add over blocks. - field_dot = accumulate_or_assign(field_dot, field_dot_pb) - if has_apar_dot: - apar_dot = accumulate_or_assign(apar_dot, apar_dot_pb) - fdot = accumulate_or_assign(fdot, fdot_pb) - src = accumulate_or_assign(src, src_pb) - bflux_tot = accumulate_or_assign(bflux_tot, bflux_tot_pb) - - - # List of handles to lines plotted, and plot a reference line at y=0. - hpl1a = list() - hpl1a.append(ax1a.plot([-1.0,1.0], [0.0,0.0], color='grey', linestyle=':', linewidth=1)) - - if not kwargs["relative_error"]: - # Plot every term in the particle balance. - - src[0] = 0.0 # Set source=0 at t=0 since we don't have fdot and bflux then. - - # Compute the error. - if has_apar_dot: - mom_err = src - bflux_tot - (fdot - field_dot - apar_dot) - else: - mom_err = src - bflux_tot - (fdot - field_dot) - - # Plot. - legend_strings = list() - - if has_src: - hpl1a.append(ax1a.plot(time_src, absy_func(src), linestyle=line_styles[2])) - legend_strings.append(r'$\mathcal{S}$') - - if has_bflux: - hpl1a.append(ax1a.plot(time_bflux_tot, absy_func(-bflux_tot), linestyle=line_styles[1])) - legend_strings.append(r'$-\int_{\partial \Omega}\mathrm{d}\mathbf{S}\cdot\mathbf{\dot{R}}f$') - - if has_field_dot: - hpl1a.append(ax1a.plot(time_field_dot, absy_func(-field_dot), linestyle=':', marker='+',markevery=8)) - legend_strings.append(r'$-\dot{\phi}$') - - if has_apar_dot: - hpl1a.append(ax1a.plot(time_apar_dot, absy_func(-apar_dot), linestyle=':', marker='+',markevery=8)) - legend_strings.append(r'$-\dot{A}_{\parallel}$') - - hpl1a.append(ax1a.plot(time_fdot, absy_func(-fdot), linestyle=line_styles[0])) - legend_strings.append(r'$-\dot{f}$') - - hpl1a.append(ax1a.plot(time_fdot, absy_func(mom_err), linestyle=line_styles[3])) - err_str = r'$E_{\dot{\mathcal{E}}}=$' - for i in range(len(legend_strings)): - err_str = err_str + legend_strings[i] - # end - legend_strings.append(err_str) - - ylabel_string = "" - if kwargs["ylabel"]: - ylabel_string = kwargs["ylabel"] - - title_string = r'Energy balance' - if kwargs["title"]: - title_string = kwargs["title"] - - ax1a.legend([hpl1a[i][0] for i in range(1,len(hpl1a))], legend_strings, fontsize=legend_font_size, frameon=False) - - # Add datasets plotted to stack. - gdat_fdot.push(time_fdot, fdot) - data.add(gdat_fdot) - - if has_src: - gdat_src.push(time_src, src) - data.add(gdat_src) - - if has_bflux: - gdat_bflux.push(time_bflux, -bflux_tot) - data.add(gdat_bflux) - - if has_field_dot: - gdat_field_dot.push(time_field_dot, field_dot) - data.add(gdat_field_dot) - - if has_apar_dot: - gdat_apar_dot.push(time_apar_dot, apar_dot) - data.add(gdat_apar_dot) - - gdat_err = GData(tag="err", label="err", ctx=gdat_fdot.ctx) - gdat_err.push(time_fdot, mom_err) - data.add(gdat_err) - - else: - # Plot the relative error. - - # Read the time step. - if kwargs["dt_file"]: - dt_file = kwargs["path"] + kwargs["dt_file"] - else: - dt_file = file_path_prefix.replace("_b*","") + 'dt.gkyl' - - _, time_dt, dt, gdat = read_gfile_if_present(dt_file) - gdat_rel_err = GData(tag="rel_err", label="rel_err", ctx=gdat.ctx) - - field = None - apar = None - distf = None - for bI in range(num_blocks): - - block_path_prefix = file_path_prefix.replace("*",str(bI)) - - # Load field energy data. - if kwargs["field_file"]: - field_file = kwargs["path"] + kwargs["field_file"].replace("*",str(bI)) - else: - field_file = block_path_prefix + 'field_energy.gkyl' - - has_field, time_field, field_pb, gdat = read_gfile_if_present(field_file) - - # Load apar energy data (optional, may not exist in electrostatic simulations). - if kwargs["apar_file"]: - apar_file = kwargs["path"] + kwargs["apar_file"].replace("*",str(bI)) - else: - apar_file = block_path_prefix + 'apar_energy.gkyl' - - has_apar, time_apar, apar_pb, gdat = read_gfile_if_present(apar_file) - - distf_pb = None - for sI in range(len(species_names)): - spec_nm = species_names[sI] - - # Load integrated moments and time step. - if kwargs["f_file"]: - f_file = (kwargs["path"] + kwargs["f_file"].replace("*",str(bI),1)).replace("*",spec_nm) - else: - f_file = block_path_prefix + spec_nm + '_integrated_moms.gkyl' - - _, time_distf, distf_ps, _ = read_gfile_if_present(f_file) - - #[ Select the Hamiltonian moment. - distf_ps = distf_ps[:,2] - - # Add over species. - distf_pb = accumulate_or_assign(distf_pb, distf_ps) - - #[ Add over blocks. - field = accumulate_or_assign(field, field_pb) - if has_apar: - apar = accumulate_or_assign(apar, apar_pb) - distf = accumulate_or_assign(distf, distf_pb) - - # Remove the t=0 data point. - field = field[1:] - field_dot = field_dot[1:] - if has_apar: - apar = apar[1:] - apar_dot = apar_dot[1:] - fdot = fdot[1:] - src = src[1:] - bflux_tot = bflux_tot[1:] - distf = distf[1:] - - # Compute the relative error. - if has_apar: - mom_err = src - bflux_tot - (fdot - field_dot - apar_dot) - mom_err_norm = mom_err*dt/(distf-field-apar) - else: - mom_err = src - bflux_tot - (fdot - field_dot) - mom_err_norm = mom_err*dt/(distf-field) - - # Plot. - hpl1a.append(ax1a.plot(time_dt, absy_func(mom_err_norm))) - - ylabel_string = r'$E_{\dot{\mathcal{E}}}~\Delta t/\mathcal{E}$' - if kwargs["ylabel"]: - ylabel_string = kwargs["ylabel"] - - title_string = r'Relative error in energy conservation' - if kwargs["title"]: - title_string = kwargs["title"] - - # Add datasets plotted to stack. - gdat_rel_err.push(time_dt, mom_err_norm) - data.add(gdat_rel_err) - - if kwargs["logy"]: - ax1a.set_yscale("log") - - if kwargs["absy"] and ylabel_string != '': - ylabel_string = r'|'+ylabel_string+r'|' - - ax1a.set_xlabel(kwargs["xlabel"],fontsize=xy_label_font_size) - ax1a.set_ylabel(ylabel_string,fontsize=xy_label_font_size) - ax1a.set_title(title_string,fontsize=title_font_size) - ax1a.set_xlim( time_fdot[0], time_fdot[-1] ) - set_tick_font_size(ax1a,tick_font_size) - - if kwargs["saveas"]: - plt.savefig(kwargs["saveas"]) - else: - plt.show() - - verb_print(ctx, "Finishing particle balance.") diff --git a/src/postgkyl/commands/gk_fluxsurf.py b/src/postgkyl/commands/gk_fluxsurf.py deleted file mode 100644 index b3cd1161..00000000 --- a/src/postgkyl/commands/gk_fluxsurf.py +++ /dev/null @@ -1,155 +0,0 @@ -import os - -import click -import numpy as np -from scipy.interpolate import PchipInterpolator - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.utils.gk_utils as gku - -# We don't import _binormal_project anymore, we handle it natively and much faster here. -from .gk_rz import ( - _file_prefix, _nodes_geometry, _mapc2p_geometry, - _interp, _centers, _sample -) - -@click.command(name="gk-fluxsurf") -@click.option("--mapc2p", "-m", default=None, type=click.STRING, - help="Use a modal mapc2p file as the geometry source instead of the default nodes file.") -@click.option("--nodes", "-n", default=None, type=click.STRING, - help="Path to a nodal geometry file, overriding the default lookup.") -@click.option("--x-idx", "-x", default=0, type=click.INT, - help="The cell index in the radial (x) direction representing the flux surface. Default 0.") -@click.option("--nphi", default=128, type=click.INT, - help="Number of toroidal angle (phi) slices. Increased default to 128 for smoother diagonal field lines.") -@click.option("--nz-interp", default=8, type=click.INT, - help="Parallel (z) up-sampling factor used to smooth the projected 3D surfaces. Default 8.") -@click.option("--use", "-u", default=None, - help="Specify tag of datasets to process from the stack.") -@click.option("--tag", "-t", default="fluxsurf", type=click.STRING, - help="Tag for output datasets.") -@click.option("--label", "-l", default=None, type=click.STRING, - help="Custom label for the result.") -@click.pass_context -def gk_fluxsurf(ctx, **kwargs): - """ - Gyrokinetics: Extract a 2D theta-phi flux surface. - - This command extracts data along a specific radial flux surface (constant x) - for 3D field-aligned data. It achieves this by performing a binormal - projection over a scan of toroidal angles (phi), creating a 2D grid of - phi vs z (where z maps along the poloidal/theta direction). - """ - data = ctx.obj["data"] - - first_data = next(data.iterator(kwargs["use"]), None) - if first_data is None: - return - - if first_data.get_num_dims() < 3: - ctx.fail("gk-fluxsurf requires 3D data to scan over toroidal angle (phi).") - - prefix = _file_prefix(getattr(first_data, "_file_name", None)) - - mapc2p_opt = kwargs["mapc2p"] - nodes_opt = kwargs["nodes"] - if mapc2p_opt is not None and nodes_opt is not None: - raise click.ClickException("Pass either --mapc2p or --nodes, not both.") - - if nodes_opt is not None: - geo_path, geo_reader = nodes_opt, _nodes_geometry - elif mapc2p_opt is not None: - geo_path = mapc2p_opt if mapc2p_opt else ( - prefix + "-geo_int_mapc2p.gkyl" if prefix is not None else None) - geo_reader = _mapc2p_geometry - elif prefix is not None: - geo_path, geo_reader = prefix + "-geo_int_nodes.gkyl", _nodes_geometry - if not os.path.exists(geo_path): - geo_path, geo_reader = prefix + "-geo_int_mapc2p.gkyl", _mapc2p_geometry - else: - geo_path, geo_reader = None, None - - if geo_path is None or not os.path.exists(geo_path): - raise click.ClickException( - "Could not find a geometry file; pass it with -n/--nodes or -m/--mapc2p.") - - x_idx = kwargs["x_idx"] - nphi = kwargs["nphi"] - nz_interp = max(1, kwargs["nz_interp"]) - - verb_print(ctx, f"Extracting theta-phi flux surface at x-index {x_idx} using geometry {geo_path}") - - # Load fine computational grid - fine_grid, _ = _interp(first_data) - xc, yc, zc = _centers(fine_grid) - Nz = zc.size - - # Load and sample the geometry to get continuous physical phi values - gx_gy_gz, majorR, vertZ, phi = geo_reader(geo_path) - gx, gy, gz = gx_gy_gz - - # Up-sample z for a smooth parallel mapping - zf_edges = np.linspace(fine_grid[2][0], fine_grid[2][-1], nz_interp * Nz + 1) - zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) - - # Unwrap toroidal angle coordinates to track continuous winding - phi = np.unwrap(np.unwrap(np.unwrap(phi, axis=2), axis=1), axis=0) - phi_grid = _sample(phi, [gx, gy, gz], [xc, yc, zf]) - - # Array of toroidal angles to scan over (standard 0 to 2*pi) - phi_tor_list = np.linspace(0, 2*np.pi, nphi, endpoint=False) - - loaded_count = 0 - for dat in data.iterator(kwargs["use"]): - _, vals = _interp(dat) - - if x_idx >= vals.shape[0] or x_idx < 0: - ctx.fail(f"Requested x-index {x_idx} is out of bounds for data with Nx={vals.shape[0]}") - - # 1. Up-sample the 3D field along the z (parallel) direction - vals_zf = PchipInterpolator(zc, vals, axis=-1, extrapolate=True)(zf) - - # 2. Extract ONLY the target radial index to save massive amounts of compute - vals_2d = vals_zf[x_idx, :, :] # Shape: (Ny, len(zf)) - phi_2d = phi_grid[x_idx, :, :] # Shape: (Ny, len(zf)) - Ny = vals_2d.shape[0] - - # 3. Allocate the 2D array for our theta-phi output grid - flux_surf_data = np.empty((nphi, len(zf))) - - # 4. Vectorized Projection: Loop only over z, evaluate all phi_tor simultaneously - for iz in range(len(zf)): - phi_y = phi_2d[:, iz] - val_y = vals_2d[:, iz] - - # Toroidal angle subtended by one full (periodic) binormal box. - box = np.mean(np.diff(phi_y)) * Ny - - # Extend domain for periodic interpolation - phi_ext = np.concatenate([phi_y - box, phi_y, phi_y + box]) - val_ext = np.concatenate([val_y, val_y, val_y]) - - # Sort to prepare for numpy interpolation - order = np.argsort(phi_ext) - phi_ext_sorted = phi_ext[order] - val_ext_sorted = val_ext[order] - - # Fold ALL requested phi angles into the local domain at once - pt_array = phi_y[0] + np.mod(phi_tor_list - phi_y[0], box) - - # Interpolate and assign the entire column of toroidal angles in one call - flux_surf_data[:, iz] = np.interp(pt_array, phi_ext_sorted, val_ext_sorted) - - out = GData(tag=kwargs["tag"], label=kwargs["label"], ctx=dat.ctx) - - # Push data back to the stack: Dimensions are now [phi, z] - out.push([phi_tor_list, zf], flux_surf_data[..., np.newaxis]) - data.add(out) - dat.deactivate() - loaded_count += 1 - - if loaded_count > 1: - data.set_unique_labels() - - verb_print(ctx, "Finishing flux surface extraction.") \ No newline at end of file diff --git a/src/postgkyl/commands/gk_load_quantity.py b/src/postgkyl/commands/gk_load_quantity.py deleted file mode 100644 index 558690dc..00000000 --- a/src/postgkyl/commands/gk_load_quantity.py +++ /dev/null @@ -1,173 +0,0 @@ -import re - -import os - -import click - -from postgkyl.utils.gk_quantities.registry import gk_quant_registry -from postgkyl.utils import verb_print - -@click.command(name="gk-load-quantity") -@click.option("--quantity", "-q", required=False, type=click.STRING, - help="Quantity to plot.") -@click.option("--qlist", is_flag=True, default=False, - help="List accepted quantities.") -@click.option("--name", "-n", required=False, type=click.STRING, - help="Simulation name prefix (e.g. gk_sheath_2x2v_p1).") -@click.option("--species", "-s", required=False, type=click.STRING, - help="Species name (e.g. ion or elc).") -@click.option("--frame", "-f", required=False, type=click.STRING, - help="Frame number, comma-separated list, or range 'start:stop[:step]'. " - "Use ':' for all available frames.") -@click.option("--path", "-p", default="./", type=click.STRING, - help="Directory containing the simulation files.") -@click.option("--tag", "-t", default="default", type=click.STRING, - help="Tag for the output dataset.") -@click.option("--label", "-l", default=None, type=click.STRING, - help="Label override for the output dataset.") -@click.option("--extra", "-e", default=None, type=click.STRING, - help="Extra comma-separated key=value pairs of extra commands, e.g. dir=1,mass=0.1. " - "A key may be given one value per species as a comma-separated array, e.g. " - "mass=me,mi1,mi2 alongside --species elc,ion1,ion2. Purpose depends on -q.") -@click.pass_context -def gk_load_quantity(ctx, **kwargs): - """ - Gyrokinetics: load a pre-named quantity from simulation output files. - - \b - For a list of accepted quantities use: - pgkyl gk-load-quantity --qlist - - \b - Command line example: - pgkyl gk-load-quantity den -s ion -n gk_sheath_2x2v_p1 -f 9 interp plot - - \b - Script example: - from postgkyl.commands.gk_load_quantity import load_gk_quantity - gdat = load_gk_quantity("n", "ion", "gk_sheath_2x2v_p1", frame=9) - """ - - if kwargs['qlist']: - # Print accepted quantities and exit. - valid = gk_quant_registry.list() - print(f"Available quantities: {', '.join(valid)}.") - return - - data = ctx.obj["data"] - verb_print(ctx, f"Loading quantity {kwargs['quantity']} for {kwargs['name']}") - - if not gk_quant_registry.has(kwargs['quantity']): - valid = gk_quant_registry.list() - raise ValueError(f"Unknown quantity '{kwargs['quantity']}'. " - f"Available quantities: {', '.join(valid)}.") - - gkquant = gk_quant_registry.get(kwargs['quantity']) - - # Parse --extra into a dict, auto-converting numeric values. - user_extra = {} - if kwargs.get('extra'): - for pair in re.split(r"[,\s]+(?=[^\s,=]+=)", kwargs['extra'].strip()): - key, _, val = pair.partition("=") - vals = [] - for v in val.split(","): - v = v.strip() - if not v: - continue - try: - v = int(v) - except ValueError: - try: - v = float(v) - except ValueError: - pass - vals.append(v) - # A single value stays a scalar and applies to every species. - user_extra[key.strip()] = vals[0] if len(vals) == 1 else vals - - path = kwargs['path'].rstrip("/") + "/" - - # Create species list. - species_inp = kwargs['species'] - species_list = [s.strip() for s in species_inp.split(",")] if species_inp else [None] - - verb_print(ctx, f"Species: {species_list}") - - if gkquant.is_multi_species: - # Combine every species into a single dataset (e.g. the sound speed), so it is fetched - # once for the whole species list instead of once per species. - if species_list == [None]: - raise ValueError(f"Quantity '{gkquant.name}' combines several species, so it needs " - f"a species list, e.g. --species elc,ion.") - - src_combo_idx, frames = gkquant.get_avail_source_multi(path, kwargs['name'], species_list, kwargs['frame']) - - verb_print(ctx, f" {species_list}: will compute {gkquant.name} using source {src_combo_idx}, frames {frames}") - - for frame in frames: - # Load required datasets (sources) for every species and compute the quantity. - out = gkquant.fetch_multi(path, kwargs['name'], species_list, frame, src_combo_idx, **user_extra) - - out_label = kwargs['label'] if kwargs['label'] is not None else gkquant.get_label() - if len(frames) > 1: - out_label += f" f{frame}" - - out.set_label(out_label) - out.set_tag(kwargs['tag']) - - data.add(out) # Push data to stack. - - verb_print(ctx, f"Finished loading '{gkquant.name}'") - return - - for species_idx, species in enumerate(species_list): - # Determine which source combination and frames to use for this species. - src_combo_idx, frames = gkquant.get_avail_source(path, kwargs['name'], species, kwargs['frame']) - - verb_print(ctx, f" {species}: will compute {gkquant.name} using source {src_combo_idx}, frames {frames}") - - # Tells the fetch functions which entry of a per-species '--extra' array - # (e.g. 'mass=1,2,3') applies to the species being computed. - species_extra = dict(user_extra, species_idx=species_idx) - - for frame in frames: - - # Load required datasets (sources) and compute the quantity. - out = gkquant.fetch(path, kwargs['name'], species, frame, src_combo_idx, **species_extra) - - # stamp a filename so that commands such as gk-rz can locate sibling files (e.g. the geometry) from the stack. - tail = f"{species}_{gkquant.name}" if species else gkquant.name - out._file_name = os.path.join(path, f"{kwargs['name']}-{tail}_{frame}.gkyl") - - # Set label. - default_label = gkquant.get_label(species=species, direction=user_extra.get("dir", None)) - - out_label = '' - if kwargs['label'] is not None: - out_label = kwargs['label'] - if len(species_list) > 1: - out_label += f" {species}" - # end - else: - out_label = default_label - - if len(frames) > 1: - out_label += f" f{frame}" - # end - - out.set_label(out_label) - - # Set tag. - out_tag = kwargs['tag'] - if len(species_list) > 1: - out_tag += f"_{species}" - # end - - out.set_tag(out_tag) - - data.add(out) # Push data to stack. - # end frame loop - # end species loop - - verb_print(ctx, f"Finished loading '{gkquant.name}'") - diff --git a/src/postgkyl/commands/gk_nodes.py b/src/postgkyl/commands/gk_nodes.py deleted file mode 100644 index ab1cd453..00000000 --- a/src/postgkyl/commands/gk_nodes.py +++ /dev/null @@ -1,316 +0,0 @@ -import click -import numpy as np -import matplotlib.pyplot as plt -import os -import glob -from matplotlib.collections import LineCollection -from itertools import cycle -from typing import Tuple - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.utils.gk_utils as gku -import postgkyl.utils.gkeyll_enums as gkenums - - - -def str_append_multib_suffix_mb(str_in, suffix, bidx): - # Append the suffix to the input string str_in and format it with the block - # index bidx. - return str_in + suffix % bidx - -def str_append_multib_suffix_sb(str_in, suffix, bidx): - # Just return the input string. - return str_in - -@click.command() -@click.option("--name", "-n", required=True, type=click.STRING, default=None, - help="Simulation name (also the file prefix, e.g. gk_sheath_1x2v_p1).") -@click.option("--path", "-p", type=click.STRING, default='./.', - help="Path to simulation data.") -@click.option("--multib", "-m", type=click.STRING, is_flag=False, flag_value="-1", default="-10", - help="Multiblock. Optional: pass block indices as comma-separated list or slice (start:stop:step). If no indices are given, all blocks are used.") -@click.option("--nodes_file", type=click.STRING, default=None, - help="Grid nodes (.gkyl format).") -@click.option("--psi_file", type=click.STRING, default=None, - help="Poloidal flux (.gkyl format).") -@click.option("--wall_file", type=click.STRING, default=None, - help="Vacuum vessel wall (.csv format).") -@click.option("--contour", "-c", is_flag=True, help="Plot contours of psi.") -@click.option("--clevels", type=click.STRING, - help="Specify levels for contours: comma-separated level values or start:end:nlevels.") -@click.option("--cnlevels", type=click.INT, default=11, help="Specify the number of levels for contours.") -@click.option("--fix_aspect", "-a", "fixaspect", is_flag=True, - help="Enforce the same scaling on both axes.") -@click.option("--xlim", default=None, type=click.STRING, - help="Set limits for the x-coordinate (lower,upper)") -@click.option("--ylim", default=None, type=click.STRING, - help="Set limits for the y-coordinate (lower,upper).") -@click.option("--xlabel", type=click.STRING, default="R (m)", - help="Label for the x axis.") -@click.option("--ylabel", type=click.STRING, default="Z (m)", - help="Label for the y axis.") -@click.option("--zlabel", type=click.STRING, default=r"$\psi$", - help="Label for the color bar.") -@click.option("--title", type=click.STRING, default=None, - help="Title for the figure.") -@click.option("--indent_left", type=click.FLOAT, default=0.0, - help="A number in the [-0.11,0.88] range by which to shift the left boundary of the plot.") -@click.option("--add_width", type=click.FLOAT, default=0.0, - help="A number in the [-0.86,0.13] range by which to increase the width the plot.") -@click.option("--multib_unicolor", is_flag=True, default=False, help="Use one color for all blocks.") -@click.option("--saveas", type=click.STRING, default=None, - help="Name of figure file.") -@click.option("--no_show", is_flag=True, default=False, - help="Suppreses showing the figure.") -@click.pass_context -def gk_nodes(ctx, **kwargs): - """ - \b - Gyrokinetics: Plot nodes of the grid, with an option to overlay - contours of the poloidal flux. - - \b - The default assumes these are in the current directory. - Alternatively, the full path to each file can be specified. - - \b - If simulation is multiblock, and you wish to specify files manually: - 1) Pass * for the block index. - 2) Use --multib/-m to specify desired blocks (or ommit to use all). - - NOTE: this command cannot be combined with other postgkyl commands. - """ - - data = ctx.obj["data"] # Data stack. - ctx.obj["plot_handles"] = {} # Handles to objects in plot. - handles = ctx.obj["plot_handles"] - - verb_print(ctx, "Plotting nodes for " + kwargs["name"]) - - kwargs["path"] = kwargs["path"] + '/' # For safety. - - # File name root including path. - if kwargs["multib"] == "-10": - file_path_prefix = kwargs["path"] + kwargs["name"] + '-' # Single block. - else: - file_path_prefix = kwargs["path"] + kwargs["name"] + '_b*-' # Multi block. - # end - - # File with nodes to plot. - if kwargs["nodes_file"]: - if kwargs["nodes_file"][0] == "/": - # Absolute path included in node file. Don't append path. - nodes_file = kwargs["nodes_file"] - else: - nodes_file = kwargs["path"] + kwargs["nodes_file"] - #end - else: - nodes_file = file_path_prefix + 'nodes.gkyl' - # end - - # Determine number of blocks. - blocks = gku.get_block_indices(kwargs["multib"], nodes_file) - num_blocks = len(blocks) - # Tag for dataset. - tag_multib_suffix = "" - str_append_multib_suffix = str_append_multib_suffix_sb - if num_blocks > 1: - tag_multib_suffix = "_b%d" - str_append_multib_suffix = str_append_multib_suffix_mb - - block_path_prefix = file_path_prefix - - # Loop through blocks to find extrema. - majorR_ex = [1e9, -1e9] - vertZ_ex = [1e9, -1e9] - for bI in blocks: - block_path_prefix = file_path_prefix.replace("*",str(bI)) - - # Load nodes. - grid, nodes, gdat = gku.read_gfile(nodes_file.replace("*",str(bI))) - - is_mapc2p = gku.is_gdata_geo_mapc2p(gdat) - majorR, vertZ = gku.nodes_to_RZ(nodes, is_mapc2p) # Major radius and vertical location. - - majorR_ex = [min([majorR_ex[0],np.amin(majorR)]), max([majorR_ex[1],np.amax(majorR)])] - vertZ_ex = [min([vertZ_ex[0],np.amin(vertZ)]), max([vertZ_ex[1],np.amax(vertZ)])] - # end - - # Create figure. - Rmin, Rmax = majorR_ex[0], majorR_ex[1] - Zmin, Zmax = vertZ_ex[0], vertZ_ex[1] - lengthR, lengthZ = Rmax-Rmin, Zmax-Zmin - aspect_ratio = lengthR/lengthZ - - ax_pos = [0.82-(8.36*aspect_ratio)/(8.36*aspect_ratio+2.5)+kwargs["indent_left"], 0.08, - (8.36*aspect_ratio)/(8.36*aspect_ratio+2.5)+kwargs["add_width"], 0.88] - cax_pos = [ax_pos[0]+ax_pos[2]+0.01, ax_pos[1], 0.02, ax_pos[3]]; - fig_prop = (8.36*aspect_ratio+2.5, 8.36+1.14) - fig_h = plt.figure(figsize=fig_prop) - ax_h = fig_h.add_axes(ax_pos) - - # Store figure handles in case script mode wishes to modify them. - handles["figure"] = fig_h - handles["axis"] = ax_h - - # Color cycler for plotting each block in a different color. - color_list = plt.rcParams['axes.prop_cycle'].by_key()['color'] - block_colors = cycle(color_list) - if kwargs["multib_unicolor"]: - block_colors = cycle([color_list[0]]) - # end - - # Loop through blocks to plot. - pl_nodes_h = list() - pl_edges_h = list() - for bI in blocks: - - block_path_prefix = file_path_prefix.replace("*",str(bI)) - # Load nodes. - grid, nodes, gdat = gku.read_gfile(nodes_file.replace("*",str(bI))) - - is_mapc2p = gku.is_gdata_geo_mapc2p(gdat) - majorR, vertZ = gku.nodes_to_RZ(nodes, is_mapc2p) # Major radius and vertical location. - - # Plot each node. - pl_nodes_h.append(ax_h.plot(majorR,vertZ,marker=".", color="k", linestyle="none")) - - cdim = np.size(np.shape(nodes))-1 - # Connect nodes with line segments. - cell_color = next(block_colors) - if (cdim == 1): - pl_edges_h.append(ax_h.plot(majorR,vertZ,color=cell_color, linestyle="-")) - else: - segs_constx = np.stack((majorR,vertZ), axis=2) - segs_consty = segs_constx.transpose(1,0,2) - pl_edges_h.append(ax_h.add_collection(LineCollection(segs_constx, color=cell_color))) - pl_edges_h.append(ax_h.add_collection(LineCollection(segs_consty, color=cell_color))) - - # Add datasets plotted to stack. - gdat_edges = GData(tag=str_append_multib_suffix("edges",tag_multib_suffix,bI), ctx=gdat.ctx) - gdat_edges.push(segs_constx, segs_consty) - data.add(gdat_edges) - # end - - # Add datasets plotted to stack. - gdat_nodes = GData(tag=str_append_multib_suffix("nodes",tag_multib_suffix,bI), ctx=gdat.ctx) - gdat_nodes.push(majorR, vertZ) - data.add(gdat_nodes) - # end - - handles["nodes"] = pl_nodes_h - handles["edges"] = pl_edges_h - - if kwargs["psi_file"]: - if kwargs["psi_file"][0] == "/": - # Absolute path included in node file. Don't append path. - psi_file = kwargs["psi_file"] - else: - psi_file = kwargs["path"] + kwargs["psi_file"] - #end - - colorbar = True - # Plot poloidal flux. - psi_grid, psi_values, gdat = gku.read_interp_gfile(psi_file, 2, 'mt') - # Convert nodal to cell center coordinates. - psi_grid_cc = list() - for d in range(len(psi_grid)): - psi_grid_cc.append(0.5*(psi_grid[d][:-1] + psi_grid[d][1:])) - # end - - if kwargs["contour"]: - # Contour plot. - if kwargs["clevels"]: - if ":" in kwargs["clevels"]: - s = clevels.split(":") - psi_clevels = np.linspace(float(s[0]), float(s[1]), int(s[2])) - else: - psi_clevels = np.array(kwargs["clevels"].split(",")) - # Filter out empty elements - psi_clevels = np.array(list(filter(None, psi_clevels))) - # end - else: - psi_clevels = kwargs["cnlevels"] - # end - - if isinstance(psi_clevels, np.ndarray) and len(psi_clevels) == 1: - colorbar = False - # end - - pl_psi_h = ax_h.contour(psi_grid_cc[0], psi_grid_cc[1], psi_values.transpose(), psi_clevels) - - # Add colorbar. - if isinstance(psi_clevels, np.ndarray): - if np.size(psi_clevels) == 1: - colorbar = False - # end - # end - - else: - # Color plot. - pl_psi_h = ax_h.pcolormesh(psi_grid[0], psi_grid[1], psi_values.transpose(), cmap='inferno') - # end - - handles["psi"] = pl_psi_h - - if colorbar: - psi_cbar_ax_h = fig_h.add_axes(cax_pos) - psi_cbar_h = plt.colorbar(pl_psi_h, ax=ax_h, cax=psi_cbar_ax_h) - psi_cbar_h.ax.tick_params(labelsize=gku.tick_font_size) - psi_cbar_h.set_label(kwargs["zlabel"], rotation=90, labelpad=0, fontsize=gku.colorbar_label_font_size) - handles["psi_colorbar_axis"] = psi_cbar_ax_h - handles["psi_colorbar"] = psi_cbar_h - # end - - # Add datasets plotted to stack. - gdat_psi = GData(tag="psi", ctx=gdat.ctx) - if kwargs["contour"]: - gdat_psi.push(psi_grid_cc, psi_values.transpose()) - else: - gdat_psi.push(psi_grid, psi_values.transpose()) - # end - data.add(gdat_psi) - - # end - - if kwargs["wall_file"]: - # Plot the wall. - if kwargs["wall_file"][0] == "/": - # Absolute path included in node file. Don't append path. - wall_file = kwargs["wall_file"] - else: - wall_file = kwargs["path"] + kwargs["wall_file"] - #end - - wall_data = np.loadtxt(open(wall_file),delimiter=',') - wall_h = ax_h.plot(wall_data[:,0],wall_data[:,1],color="grey") - handles["wall"] = wall_h - # end - - ax_h.set_xlabel(kwargs["xlabel"],fontsize=gku.xy_label_font_size) - ax_h.set_ylabel(kwargs["ylabel"],fontsize=gku.xy_label_font_size) - ax_h.set_title(kwargs["title"],fontsize=gku.title_font_size) - if kwargs["xlim"]: - ax_h.set_xlim( float(kwargs["xlim"].split(",")[0]), float(kwargs["xlim"].split(",")[1]) ) -# else: -# ax_h.set_xlim( Rmin-0.05*lengthR, Rmax+0.05*lengthR ) - # end - - if kwargs["ylim"]: - ax_h.set_ylim( float(kwargs["ylim"].split(",")[0]), float(kwargs["ylim"].split(",")[1]) ) -# else: -# ax_h.set_ylim( Zmin-0.05*lengthZ, Zmax+0.05*lengthZ ) - # end - - gku.set_tick_font_size(ax_h,gku.tick_font_size) - - if kwargs["saveas"]: - plt.savefig(kwargs["saveas"]) - # end - - if not kwargs["no_show"]: - plt.show() - # end - - verb_print(ctx, "Finishing nodes plot.") diff --git a/src/postgkyl/commands/gk_particle_balance.py b/src/postgkyl/commands/gk_particle_balance.py deleted file mode 100644 index 80bcf65a..00000000 --- a/src/postgkyl/commands/gk_particle_balance.py +++ /dev/null @@ -1,411 +0,0 @@ -import click -import numpy as np -import matplotlib.pyplot as plt -import os -import glob - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -@click.command() -@click.option("--name", "-n", required=True, type=click.STRING, default=None, - help="Simulation name (also the file prefix, e.g. gk_sheath_1x2v_p1).") -@click.option("--species", "-s", required=True, type=click.STRING, default=None, - help="Species name.") -@click.option("--path", "-p", type=click.STRING, default='./.', - help="Path to simulation data.") -@click.option("--relative_error", "-r", is_flag=True, - help="Plot the relative error only.") -@click.option("--multib", "-m", is_flag=False, flag_value="-1", default="-10", - help="Multiblock. Optional: pass block indices as comma-separated list or slice (start:stop:step). If no indices are given, all blocks are used.") -@click.option("--fdot_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of change in f over a time step.") -@click.option("--source_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of the source(s).") -@click.option("--bflux_xlower_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through lower x boundary.") -@click.option("--bflux_ylower_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through lower y boundary.") -@click.option("--bflux_zlower_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through lower z boundary.") -@click.option("--bflux_xupper_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through upper x boundary.") -@click.option("--bflux_yupper_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through upper y boundary.") -@click.option("--bflux_zupper_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of boundary flux through upper z boundary.") -@click.option("--f_file", type=click.STRING, default=None, multiple=True, - help="Integrated moments of f.") -@click.option("--dt_file", type=click.STRING, default=None, - help="Time step.") -@click.option("--logy", is_flag=True, default=False, - help="Logarithmic scale for y axis.") -@click.option("--absy", is_flag=True, default=False, - help="Take absolute value of time traces.") -@click.option("--xlabel", type=click.STRING, default="Time (s)", - help="Label for the x axis.") -@click.option("--ylabel", type=click.STRING, default=None, - help="Label for the y axis.") -@click.option("--title", type=click.STRING, default=None, - help="Take absolute value of time traces.") -@click.option("--indent_left", type=click.FLOAT, default=0.0, - help="A number in the [-0.11,0.88] range by which to shift the left boundary of the plot.") -@click.option("--add_width", type=click.FLOAT, default=0.0, - help="A number in the [-0.86,0.13] range by which to increase the width the plot.") -@click.option("--saveas", type=click.STRING, default=None, - help="Name of figure file.") -@click.pass_context -def gk_particle_balance(ctx, **kwargs): - """ - \b - Gyrokinetics: Plot the particle balance of a given species. - Requires the following files: - ..._fdot_integrated_moms.gkyl - ..._source_integrated_moms.gkyl - ..._bflux__integrated_HamiltonianMoments.gkyl - where ... means -. - The last two files above are only needed if the simulation had - sources or non-periodic boundaries. If the relative error is - requested, these are also needed: - ..._integrated_moms.gkyl - -dt.gkyl - - \b - The default assumes these are in the current directory. - Alternatively, the full path to each file can be specified. - - \b - If simulation is multiblock, and you wish to specify files manually: - 1) Pass * for the block index. - 2) Use --multib/-m to specify desired blocks (or ommit to use all). - - NOTE: this command cannot be combined with other postgkyl commands. - """ - - # - # Hardcoded parameters and auxiliary functions. - # - max_num_blocks = 10000 - - # Labels used to identify boundary flux files. - edges = ["lower","upper"] - dirs = ["x","y","z"] - # Line styles. - line_styles = ['-','--',':','-.','None','None','None','None'] - # Font sizes. - xy_label_font_size = 17 - title_font_size = 17 - tick_font_size = 14 - legend_font_size = 14 - - # Create figure. - figProp1a = (7.5, 4.5) - ax1aPos = [0.11+kwargs["indent_left"], 0.15, 0.87+kwargs["add_width"], 0.78] - fig1a = plt.figure(figsize=figProp1a) - ax1a = fig1a.add_axes(ax1aPos) - - def set_tick_font_size(axIn,fontSizeIn): - # Set the font size of the ticks to a given size. - axIn.tick_params(axis='both',labelsize=fontSizeIn) - offset_txt = axIn.yaxis.get_offset_text() # Get the text object - offset_txt.set_size(fontSizeIn) # Set the size. - offset_txt = axIn.xaxis.get_offset_text() # Get the text object - offset_txt.set_size(fontSizeIn) # Set the size. - - def read_gfile_if_present(file_name): - # Check if a Gkeyll file exists. If it does, read it and return - # its grid, data and GData object. If it doesn't, return None. - if os.path.exists(file_name): - pgData = GData(file_name) # Read data with pgkyl. - time = pgData.get_grid() # Time stamps of the simulation. - val = pgData.get_values() # Data values. - return True, np.squeeze(time), np.squeeze(val), pgData - else: - verb_print(ctx, " -> File "+file_name+" not found. Proceeding w/o it.") - return False, None, None, None - - def parse_slice_string(value): - # Parse a 'slice()' from string, like 'start:stop:step'. - parts = value.split(':') - # Convert parts to integers, replacing empty strings with None for slice defaults - parsed_parts = [] - for p in parts: - try: - parsed_parts.append(int(p) if p else None) - except ValueError: - # Handle cases where the part might not be a number - raise ValueError(f"Invalid slice part: {p}") - # Create the slice object with the appropriate number of arguments - return slice(*parsed_parts) - - def accumulate_or_assign(target_arr, old_arr): - # Accumulates old_arr into target_arr if target_arr exists, - # otherwise assign old_arr to target_arr. - old_arr = np.asarray(old_arr) # Ensure old_arr is a numpy array. - if target_arr is None: - return old_arr.copy() - else: - target_arr += old_arr - return target_arr - - def absy_enabled(data_in): - # Take the absolute value of the data - return np.abs(data_in) - - def absy_disabled(data_in): - # Don't take the absolute value of the data - return data_in - # - # End of hardcoded parameters and auxiliary functions. - # - - data = ctx.obj["data"] # Data stack. - - verb_print(ctx, "Plotting particle balance for " + kwargs["species"] + " species.") - - absy_func = absy_disabled - if kwargs["absy"]: - absy_func = absy_enabled - - kwargs["path"] = kwargs["path"] + '/' # For safety. - - # Determine blocks to plot, number of blocks, and set file prefix. - if kwargs["multib"] == "-10": - # Single block. - file_path_prefix = kwargs["path"] + kwargs["name"] + '-' - blocks = [0] - num_blocks = 1 - else: - # Multi block. - file_path_prefix = kwargs["path"] + kwargs["name"] + '_b*-' - - if kwargs["multib"] == "-1": - # Find and use all blocks. - if kwargs["fdot_file"]: - fdot_file = kwargs["path"] + kwargs["fdot_file"] - else: - fdot_file = file_path_prefix + kwargs["species"] + '_fdot_integrated_moms.gkyl' - - fdot_file_list = glob.glob(fdot_file) - num_blocks = len(fdot_file_list) - blocks = list(range(num_blocks)) - else: - # Use specified blocks. - if ',' in kwargs["multib"]: - blocks = kwargs["multib"].split(",") - num_blocks = len(blocks) - blocks = [int(blocks[i]) for i in range(num_blocks)] - elif ':' in kwargs["multib"]: - slice_obj = parse_slice_string(kwargs["multib"]) - blocks = list(range(*slice_obj.indices(max_num_blocks))) - num_blocks = len(blocks) - - else: - raise NameError("Blocks given to --multib -m must be a comma separated list or slice.") - - block_path_prefix = file_path_prefix - - fdot = None - src = None - bflux_tot = None - for bI in range(num_blocks): - - block_path_prefix = file_path_prefix.replace("*",str(bI)) - - # Load change in species over a time step. - if kwargs["fdot_file"]: - fdot_file = kwargs["path"] + kwargs["fdot_file"].replace("*",str(bI)) - else: - fdot_file = block_path_prefix + kwargs["species"] + '_fdot_integrated_moms.gkyl' - - has_fdot, time_fdot, fdot_pb, gdat = read_gfile_if_present(fdot_file) - if not has_fdot or gdat is None: - raise FileNotFoundError(f"Required file not found: {fdot_file}") - gdat_fdot = GData(tag="fdot", label="fdot", ctx=gdat.ctx) - - # Load integrated moments of the source. - if kwargs["source_file"]: - src_file = kwargs["path"] + kwargs["source_file"].replace("*",str(bI)) - else: - src_file = block_path_prefix + kwargs["species"] + '_source_integrated_moms.gkyl' - - has_src, time_src, src_pb, gdat = read_gfile_if_present(src_file) - if has_src: - gdat_src = GData(tag="src", label="src", ctx=gdat.ctx) - - # Load particle boundary fluxes. - nbflux = 0 - time_bflux, bflux_pb = list(), list() - has_bflux = False - for d in dirs: - for e in edges: - if kwargs["bflux_"+d+e+"_file"]: - bflux_file = kwargs["path"] + kwargs["bflux_"+d+e+"_file"].replace("*",str(bI)) - else: - bflux_file = block_path_prefix + kwargs["species"] + '_bflux_'+d+e+'_integrated_HamiltonianMoments.gkyl' - - has_bflux_at_boundary, time_bflux_tmp, bflux_tmp, gdat = read_gfile_if_present(bflux_file) - if has_bflux_at_boundary: - gdat_bflux = GData(tag="bflux", label="bflux", ctx=gdat.ctx) - time_bflux.append(time_bflux_tmp) - bflux_pb.append(bflux_tmp) - has_bflux = has_bflux or has_bflux_at_boundary - nbflux += 1 - - # Select the M0 moment. - fdot_pb = fdot_pb[:,0] - if has_src: - src_pb = src_pb[:,0] - else: - src_pb = 0.0*fdot_pb - - if has_bflux: - for i in range(nbflux): - bflux_pb[i] = bflux_pb[i][:,0] - - # Add boundary fluxes of all boundaries. - if has_bflux: - time_bflux_tot = time_bflux[0] - bflux_tot_pb = bflux_pb[0] - for i in range(1,nbflux): - bflux_tot_pb += bflux_pb[i] - else: - bflux_tot_pb = 0.0*fdot_pb - - # Add over blocks. - fdot = accumulate_or_assign(fdot, fdot_pb) - src = accumulate_or_assign(src, src_pb) - bflux_tot = accumulate_or_assign(bflux_tot, bflux_tot_pb) - - - # List of handles to lines plotted, and plot a reference line at y=0. - hpl1a = list() - hpl1a.append(ax1a.plot([-1.0,1.0], [0.0,0.0], color='grey', linestyle=':', linewidth=1)) - - if not kwargs["relative_error"]: - # Plot every term in the particle balance. - - src[0] = 0.0 # Set source=0 at t=0 since we don't have fdot and bflux then. - - # Compute the error. - mom_err = src - bflux_tot - fdot - - # Plot. - legend_strings = list() - if has_src: - hpl1a.append(ax1a.plot(time_src, absy_func(src), linestyle=line_styles[2])) - legend_strings.append(r'$\mathcal{S}$') - - if has_bflux: - hpl1a.append(ax1a.plot(time_bflux_tot, absy_func(-bflux_tot), linestyle=line_styles[1])) - legend_strings.append(r'$-\int_{\partial \Omega}\mathrm{d}\mathbf{S}\cdot\mathbf{\dot{R}}f$') - - hpl1a.append(ax1a.plot(time_fdot, absy_func(-fdot), linestyle=line_styles[0])) - legend_strings.append(r'$-\dot{f}$') - - hpl1a.append(ax1a.plot(time_fdot, absy_func(mom_err), linestyle=line_styles[3])) - err_str = r'$E_{\dot{\mathcal{N}}}=$' - for i in range(len(legend_strings)): - err_str = err_str + legend_strings[i] - # end - legend_strings.append(err_str) - - ylabel_string = "" - if kwargs["ylabel"]: - ylabel_string = kwargs["ylabel"] - - title_string = r'Particle balance' - if kwargs["title"]: - title_string = kwargs["title"] - - ax1a.legend([hpl1a[i][0] for i in range(1,len(hpl1a))], legend_strings, fontsize=legend_font_size, frameon=False) - - # Add datasets plotted to stack. - gdat_fdot.push(time_fdot, fdot) - data.add(gdat_fdot) - - if has_src: - gdat_src.push(time_src, src) - data.add(gdat_src) - - if has_bflux: - gdat_bflux.push(time_bflux, -bflux_tot) - data.add(gdat_bflux) - - gdat_err = GData(tag="err", label="err", ctx=gdat_fdot.ctx) - gdat_err.push(time_fdot, mom_err) - data.add(gdat_err) - - else: - # Plot the relative error. - - if kwargs["dt_file"]: - dt_file = kwargs["path"] + kwargs["dt_file"] - else: - dt_file = file_path_prefix.replace("_b*","") + 'dt.gkyl' - - _, time_dt, dt, gdat = read_gfile_if_present(dt_file) - gdat_rel_err = GData(tag="rel_err", label="rel_err", ctx=gdat.ctx) - - distf = None - for bI in range(num_blocks): - - block_path_prefix = file_path_prefix.replace("*",str(bI)) - - # Load integrated moments and time step. - if kwargs["f_file"]: - f_file = kwargs["path"] + kwargs["f_file"].replace("*",str(bI)) - else: - f_file = block_path_prefix + kwargs["species"] + '_integrated_moms.gkyl' - - _, time_distf, distf_pb, _ = read_gfile_if_present(f_file) - - # Select the M0 moment. - distf_pb = distf_pb[:,0] - - # Add over blocks. - distf = accumulate_or_assign(distf, distf_pb) - - # Remove the t=0 data point. - fdot = fdot[1:] - src = src[1:] - bflux_tot = bflux_tot[1:] - distf = distf[1:] - - # Compute the relative error. - mom_err = src - bflux_tot - fdot - mom_err_norm = mom_err*dt/distf - - # Plot. - hpl1a.append(ax1a.plot(time_dt, absy_func(mom_err_norm))) - - ylabel_string = r'$E_{\dot{\mathcal{N}}}~\Delta t/\mathcal{N}$' - if kwargs["ylabel"]: - ylabel_string = kwargs["ylabel"] - - title_string = r'Relative error in particle conservation' - if kwargs["title"]: - title_string = kwargs["title"] - - # Add datasets plotted to stack. - gdat_rel_err.push(time_dt, mom_err_norm) - data.add(gdat_rel_err) - - if kwargs["logy"]: - ax1a.set_yscale("log") - - if kwargs["absy"] and ylabel_string != '': - ylabel_string = r'|'+ylabel_string+r'|' - - ax1a.set_xlabel(kwargs["xlabel"],fontsize=xy_label_font_size) - ax1a.set_ylabel(ylabel_string,fontsize=xy_label_font_size) - ax1a.set_title(title_string,fontsize=title_font_size) - ax1a.set_xlim( time_fdot[0], time_fdot[-1] ) - set_tick_font_size(ax1a,tick_font_size) - - if kwargs["saveas"]: - plt.savefig(kwargs["saveas"]) - else: - plt.show() - - verb_print(ctx, "Finishing particle balance.") diff --git a/src/postgkyl/commands/gk_rz.py b/src/postgkyl/commands/gk_rz.py deleted file mode 100644 index 9879e696..00000000 --- a/src/postgkyl/commands/gk_rz.py +++ /dev/null @@ -1,317 +0,0 @@ -import os - -import click -import numpy as np -from scipy.interpolate import PchipInterpolator, RegularGridInterpolator - -from postgkyl.data import GData, GInterpModal -from postgkyl.utils import verb_print -import postgkyl.utils.gk_utils as gku - - -def _file_prefix(file_name): - if not file_name: - return None - return os.path.splitext(file_name)[0].rsplit("-", 1)[0] - - -def _mapc2p_geometry(path): - """ - Interpolate a modal mapc2p (or geo R,Z,phi) file to physical R, Z, phi. - """ - gdat = GData(path) - if gku.is_gdata_geo_mapc2p(gdat): - # Cartesian X, Y, Z: R = sqrt(X^2 + Y^2), phi = atan2(Y, X). - grid, X = _interp(gdat, 0) - _, Y = _interp(gdat, 1) - _, Z = _interp(gdat, 2) - return _centers(grid), np.sqrt(X**2 + Y**2), Z, np.arctan2(Y, X) - - # Components are directly R, Z, phi. - grid, R = _interp(gdat, 0) - _, Z = _interp(gdat, 1) - phi = _interp(gdat, 2)[1] if R.ndim == 3 else None - return _centers(grid), R, Z, phi - - -def _gauss_nodes(edges): - """Coordinates of the p1 nodal points (cell center +/- h/(2*sqrt(3))) of a - 1D edge grid, i.e. where the values of a nodal geometry file live.""" - c = 0.5 * (edges[:-1] + edges[1:]) - off = np.diff(edges) / (2.0 * np.sqrt(3.0)) - return np.ravel(np.column_stack([c - off, c + off])) - - -def _nodes_geometry(path): - """ - Read a nodal geometry file (pointwise node coordinates) to physical R, Z, phi. - """ - gdat = GData(path) - vals = gdat.get_values() - # The stored grid is 2x-refined (one edge per value plus one); the values - # themselves sit at the two p1 nodes of each cell, whose edges are every - # other stored grid point. - coords = [] - for dim, g in enumerate(gdat.get_grid()): - g = np.squeeze(g) - if len(g) != vals.shape[dim] + 1 or vals.shape[dim] % 2: - raise ValueError("Unrecognized nodal geometry layout in " + path) - coords.append(_gauss_nodes(g[::2])) - - if gku.is_gdata_geo_mapc2p(gdat): - X, Y, Z = vals[..., 0], vals[..., 1], vals[..., 2] - return coords, np.sqrt(X**2 + Y**2), Z, np.arctan2(Y, X) - - R, Z = vals[..., 0], vals[..., 1] - phi = vals[..., 2] if R.ndim == 3 else None - return coords, R, Z, phi - - -def _corner_rz(path): - """R and Z from a pointwise corner geometry file ('-geo_corn_nodes.gkyl'). - - The corner file stores one value per cell corner, endpoints included, so - unlike the interior nodes/mapc2p files it covers the domain boundary (in - particular z = +/-pi, where the poloidal cross section closes). Returns - (coords, R, Z) with 'coords' the corner lattice. - """ - gdat = GData(path) - vals = gdat.get_values() - coords = [np.linspace(np.squeeze(g)[0], np.squeeze(g)[-1], n) - for g, n in zip(gdat.get_grid(), vals.shape[:-1])] - - if gku.is_gdata_geo_mapc2p(gdat): - X, Y, Z = vals[..., 0], vals[..., 1], vals[..., 2] - return coords, np.sqrt(X**2 + Y**2), Z - - return coords, vals[..., 0], vals[..., 1] - - -def _interp(gdat, comp=0): - """Interpolate component 'comp' of the DG GData object 'gdat'. - - Returns the computational grid (list of 1D node arrays) and the - interpolated values at fine cell centers. - """ - poly_order = gdat.ctx["poly_order"] - basis_type = gdat.ctx["basis_type"] - if basis_type == "serendipity": - basis_type = "ms" - - grid, vals = GInterpModal(gdat, poly_order, basis_type).interpolate(comp) - return [np.squeeze(g) for g in grid], np.squeeze(vals) - - -def _centers(nodes): - """Cell centers from a list of 1D node arrays.""" - return [0.5 * (n[:-1] + n[1:]) for n in nodes] - - -def _sample(values, src_coords, dst_coords): - """Linearly interpolate `values` onto the grid spanned by `dst_coords`.""" - mesh = np.meshgrid(*dst_coords, indexing="ij") - return RegularGridInterpolator( - tuple(src_coords), values, bounds_error=False, fill_value=None - )(tuple(mesh)) - - -def _fft_poloidal_project(vals, zc, box, wind, phi0_zf, zf, phi_tor): - """Project a 3D field-aligned dataset onto the poloidal plane at phi = phi_tor. - """ - Nx, Ny, Nz = vals.shape - fk = np.fft.rfft(vals, axis=1, norm="forward") # (Nx, K, Nz) - K = fk.shape[1] - - # Twist-and-shift reconnection: add a ghost z-cell at each domain edge (+/-pi) - # whose value is the opposite end phase-shifted by exp(i k n0 wind). - dz = zc[1] - zc[0] - z_ex = np.concatenate(([zc[0] - dz / 2], zc, [zc[-1] + dz / 2])) - fk_ex = np.zeros((Nx, K, Nz + 2), dtype=complex) - fk_ex[:, :, 1:-1] = fk - psh = (2.0 * np.pi / box) * wind # per-mode phase = n0 * wind(x) - for k in range(K): - ph = np.exp(-1j * k * psh) - fk_ex[:, k, -1] = 0.5 * (fk[:, k, -1] + ph * fk[:, k, 0]) - fk_ex[:, k, 0] = 0.5 * (fk[:, k, 0] + np.conj(ph) * fk[:, k, -1]) - - # Up-sample along z (interpolate real and imaginary parts separately). - fk_zf = (PchipInterpolator(z_ex, fk_ex.real, axis=2)(zf) - + 1j * PchipInterpolator(z_ex, fk_ex.imag, axis=2)(zf)) - - # Phase-sum: reconstruct the field where the physical toroidal angle == phi_tor. - frac = (phi_tor - phi0_zf) / box - out = np.zeros((Nx, len(zf))) - for k in range(K): - # rfft: modes 0 < k < Nyquist represent both +/-k; do not double k=0 or Nyquist. - weight = 1.0 if (k == 0 or (Ny % 2 == 0 and k == K - 1)) else 2.0 - out += weight * np.real(fk_zf[:, k, :] * np.exp(-1j * 2.0 * np.pi * k * frac)) - return out - - -@click.command() -@click.option("--mapc2p", "-m", default=None, type=click.STRING, - help="Use a modal mapc2p file as the geometry source instead of the default nodes file; " - "pass '' to look up '-geo_int_mapc2p.gkyl' from the first processed dataset's prefix.") -@click.option("--nodes", "-n", default=None, type=click.STRING, - help="Path to a nodal geometry file, overriding the default '-geo_int_nodes.gkyl' lookup.") -@click.option("--z-axis", "-z", default=0.0, type=click.FLOAT, - help="Vertical position of the magnetic axis (m), added to the geometry Z." - "mapc2p files store Z relative to the axis; pass Z_axis from the simulation input " - "file to plot in machine coordinates. Default 0.") -@click.option("--use", "-u", default=None, - help="Specify tag of datasets to process from the stack.") -@click.option("--tag", "-t", default="rz", type=click.STRING, - help="Tag for output datasets.") -@click.option("--label", "-l", default=None, type=click.STRING, - help="Custom label for the result.") -@click.option("--phi-tor", "-p", default=0.0, type=click.FLOAT, - help="Toroidal angle (radians) of the poloidal plane to project 3D data onto. Default 0.") -@click.option("--nz-interp", default=8, type=click.INT, - help="Parallel (z) up-sampling factor used to smooth the projected 3D surfaces. Default 8.") -@click.pass_context -def gk_rz(ctx, **kwargs): - """ - \b - Gyrokinetics: Interpolate DG dataset(s) and map them to the R-Z plane. - Assumes DG data (not yet interpolated) has been loaded onto the stack by a - preceding command. - - The geometry is automatically found from the prefix of the first processed - dataset: the pointwise '-geo_int_nodes.gkyl' is preferred (exact node - coordinates, robust at coarse z resolution), falling back to the modal - '-geo_int_mapc2p.gkyl'. Use '-n path' to point at a specific nodes - file, '-m path' at a specific mapc2p file, or "-m ''" to force the default - mapc2p lookup. - - For 3D (field-aligned) data the field is reconstructed on the poloidal plane at - toroidal angle --phi-tor (default 0) by interpolating along the binormal - direction, up-sampled in z (--nz-interp) for smooth surfaces. - """ - data = ctx.obj["data"] - - # Locate the geometry files from the prefix of the first processed dataset. - first_data = next(data.iterator(kwargs["use"]), None) - if first_data is None: - return - - prefix = _file_prefix(getattr(first_data, "_file_name", None)) - - # Geometry source: the pointwise nodes file by default (exact node values; - # the modal mapc2p representation loses amplitude where the toroidal winding - # is under-resolved), the modal mapc2p file on request or as fallback. - mapc2p_opt = kwargs["mapc2p"] - nodes_opt = kwargs["nodes"] - if mapc2p_opt is not None and nodes_opt is not None: - raise click.ClickException("Pass either --mapc2p or --nodes, not both.") - - if nodes_opt is not None: - geo_path, geo_reader = nodes_opt, _nodes_geometry - elif mapc2p_opt is not None: - # An empty value requests the default '-geo_int_mapc2p.gkyl'. - geo_path = mapc2p_opt if mapc2p_opt else ( - prefix + "-geo_int_mapc2p.gkyl" if prefix is not None else None) - geo_reader = _mapc2p_geometry - elif prefix is not None: - geo_path, geo_reader = prefix + "-geo_int_nodes.gkyl", _nodes_geometry - if not os.path.exists(geo_path): - geo_path, geo_reader = prefix + "-geo_int_mapc2p.gkyl", _mapc2p_geometry - else: - geo_path, geo_reader = None, None - - if geo_path is None or not os.path.exists(geo_path): - raise click.ClickException( - "Could not find a geometry file; pass it with -N/--nodes or -n/--mapc2p.") - - if first_data.get_num_dims() == 2: - # Direct map onto R-Z using mapc2p. - verb_print(ctx, "Mapping stack data to R-Z using " + geo_path) - geo_coords, majorR, vertZ, _ = geo_reader(geo_path) - vertZ = vertZ + kwargs["z_axis"] - loaded_count = 0 - for dat in data.iterator(kwargs["use"]): - field_grid, vals = _interp(dat) - # Evaluate R, Z at the field cell corners (its node arrays) so pcolormesh - # gets explicit cell edges, not non-monotonic curvilinear cell centers. - R = _sample(majorR, geo_coords, field_grid) - Z = _sample(vertZ, geo_coords, field_grid) - out = GData(tag=kwargs["tag"], label=kwargs["label"], ctx=dat.ctx) - out.push([R, Z], vals[..., np.newaxis]) - data.add(out) - dat.deactivate() - loaded_count += 1 - - if loaded_count > 1: - data.set_unique_labels() - - verb_print(ctx, "Finishing R-Z mapping.") - return - - # 3D: project onto the poloidal plane at phi_tor. - phi_tor = kwargs["phi_tor"] - nz_interp = max(1, kwargs["nz_interp"]) - - # Field cell-center coordinates (from the field's own DG grid). - fine_grid, _ = _interp(first_data) - xc, yc, zc = _centers(fine_grid) - Nx, Ny, Nz = xc.size, yc.size, zc.size - - verb_print(ctx, "3D data: projecting onto phi = %g rad using geometry %s" - % (phi_tor, geo_path)) - gx_gy_gz, majorR, vertZ, phi = geo_reader(geo_path) - if phi is None: - raise click.ClickException( - "The geometry file has no toroidal-angle component; cannot project 3D data.") - vertZ = vertZ + kwargs["z_axis"] - gx, gy, gz = gx_gy_gz - - # Physical toroidal angle on the field grid, made continuous for interpolation. - phi = np.unwrap(np.unwrap(np.unwrap(phi, axis=2), axis=1), axis=0) - phiF = _sample(phi, [gx, gy, gz], [xc, yc, zc]) - # The binormal domain spans one 1/n0 toroidal sector (wedge). - box = np.mean(np.diff(phiF[Nx // 2, :, Nz // 2])) * Ny - n0 = max(1, int(round(abs(2.0 * np.pi / box)))) - box = np.sign(box) * 2.0 * np.pi / n0 - wind = phiF[:, 0, -1] - phiF[:, 0, 0] - - # Up-sampled z: edges (zf_edges) for the plotting grid, centers (zf) for the - # field reconstruction. - xn = fine_grid[0] - zf_edges = np.linspace(fine_grid[2][0], fine_grid[2][-1], nz_interp * Nz + 1) - zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) - # Toroidal angle at the binormal origin along (x, zf), used by the phase-sum. - phi0_zf = np.array([np.interp(zf, zc, phiF[ix, 0, :]) for ix in range(Nx)]) - - # R, Z (y-independent) on the (x, z) plane. - R2d, Z2d, gz_rz = majorR[:, 0, :], vertZ[:, 0, :], gz - corn_path = prefix + "-geo_corn_nodes.gkyl" if prefix is not None else None - if corn_path is not None and os.path.exists(corn_path): - verb_print(ctx, "Closing the theta = +/-pi ends with " + corn_path) - ccoords, cornR, cornZ = _corner_rz(corn_path) - cx, cz = ccoords[0], ccoords[2] - cornR, cornZ = cornR[:, 0, :], cornZ[:, 0, :] + kwargs["z_axis"] - R2d = np.concatenate([ - np.interp(gx, cx, cornR[:, 0])[:, None], R2d, - np.interp(gx, cx, cornR[:, -1])[:, None]], axis=1) - Z2d = np.concatenate([ - np.interp(gx, cx, cornZ[:, 0])[:, None], Z2d, - np.interp(gx, cx, cornZ[:, -1])[:, None]], axis=1) - gz_rz = np.concatenate([[cz[0]], gz, [cz[-1]]]) - - # Interpolate the geometry. - Rrz = _sample(R2d, [gx, gz_rz], [xn, zf_edges]) - Zrz = _sample(Z2d, [gx, gz_rz], [xn, zf_edges]) - - loaded_count = 0 - for dat in data.iterator(kwargs["use"]): - _, vals = _interp(dat) - proj = _fft_poloidal_project(vals, zc, box, wind, phi0_zf, zf, phi_tor) - out = GData(tag=kwargs["tag"], label=kwargs["label"], ctx=dat.ctx) - out.push([Rrz, Zrz], proj[..., np.newaxis]) - data.add(out) - dat.deactivate() - loaded_count += 1 - - if loaded_count > 1: - data.set_unique_labels() - - verb_print(ctx, "Finishing R-Z mapping.") diff --git a/src/postgkyl/commands/gkyl_pkpm.py b/src/postgkyl/commands/gkyl_pkpm.py deleted file mode 100644 index 85ca2bdd..00000000 --- a/src/postgkyl/commands/gkyl_pkpm.py +++ /dev/null @@ -1,42 +0,0 @@ -import click - -from postgkyl.data import GData, GInterpModal -from postgkyl.utils import verb_print -import postgkyl.tools.laguerre_compose -import postgkyl.tools.transform_frame - - -@click.command() -@click.option("--name", "-n", type=click.STRING, prompt=True, help="Set the root name for files.") -@click.option("--species", "-s", type=click.STRING, prompt=True, help="Set species name.") -@click.option("--idx", "-i", type=click.STRING, prompt=True, help="Set the file number.") -@click.option("--poly_order", "-p", type=click.INT, prompt=True, help="Set the polynomial order.") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def pkpm(ctx, **kwargs): - """Shortcut to load Gkeyll PKPM data, interpolate, and transform.""" - verb_print(ctx, "Starting Gkyl PKPM") - data = ctx.obj["data"] - - gf = GData(f"{kwargs['name'],:s}-{kwargs['species']:s}_{kwargs['idx']:s}.gkyl") - gvars = GData(f"{kwargs['name']:s}-{kwargs['species']:s}_pkpm_vars_{kwargs['idx']:s}.gkyl") - - num_dims = gf.get_num_dims() - c_dim = num_dims - 1 - - dg = GInterpModal(gf, kwargs["poly_order"], "pkpmhyb") - dg.interpolate((0, 1), overwrite=True) - - dg = GInterpModal(gvars, kwargs["poly_order"], "ms") - grid_and_T_m = dg.interpolate(3) - grid_and_us = dg.interpolate((0, 1, 2)) - - postgkyl.tools.laguerre_compose(gf, grid_and_T_m, gf) - postgkyl.tools.transform_frame(gf, grid_and_us, c_dim, gf) - - gf.set_tag(kwargs["tag"]) - gf.set_label(kwargs["label"]) - data.add(gf) - - verb_print(ctx, "Finishing Gkyl PKPM") diff --git a/src/postgkyl/commands/grid.py b/src/postgkyl/commands/grid.py deleted file mode 100644 index 863ac6b5..00000000 --- a/src/postgkyl/commands/grid.py +++ /dev/null @@ -1,54 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", type=click.STRING, help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.option("--read", "-r", type=click.BOOL, help="Read from general interpolation file.") -@click.pass_context -def grid(ctx, **kwargs): - """Create a dataset out of a grid""" - verb_print(ctx, "Starting grid") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - grid_in = dat.get_grid() - num_dims = dat.get_num_dims() - num_cells = dat.get_num_cells() - grid_out = [] - for nc in num_cells: - grid_out.append(np.arange(nc+2)) - # end - - shape = np.copy(num_cells) + 1 - shape = np.append(shape, num_dims) - values = np.zeros(shape) - - if num_dims == 1: - values[..., 0] = grid_in[0] - elif len(grid_in[0].shape) == 1: # uniform mesh or vel c2p mapping - temp = np.meshgrid(*grid_in, indexing="ij") - for d, t in enumerate(temp): - values[..., d] = t - # end - else: # c2p mapping - for d, t in enumerate(grid_in): - values[..., d] = t - # end - # end - - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - out.push(grid_out, values) - data.add(out) - else: - dat.push(grid_out, values) - # end - # end - verb_print(ctx, "Finishing grid") diff --git a/src/postgkyl/commands/growth.py b/src/postgkyl/commands/growth.py deleted file mode 100644 index b3abdaef..00000000 --- a/src/postgkyl/commands/growth.py +++ /dev/null @@ -1,104 +0,0 @@ -import click -import matplotlib.pyplot as plt -import numpy as np -import os - -from postgkyl.data import GData -import postgkyl.tools -from postgkyl.utils import verb_print - - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("-g", "--guess", help="Specify comma-separated initial guess.") -@click.option("--minn", type=click.INT, help="Set minimal number of points to fit.") -@click.option("-d", "--dataset", is_flag=True, help="Create a new dataset with fitted exponential.") -@click.option("-i", "--instantaneous", is_flag=True, help="Plot instantaneous growth rate vs time.") -@click.option("--dir", type=click.INT, help="Choose direction for multi-D data.") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def growth(ctx, **kwargs): - """Attempts to compute growth rate (i.e. fit e^(2x)) from DynVector data. - - the DynVector is typically an integrated quantity like electric or magnetic field - energy. - """ - verb_print(ctx, "Starting growth") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - time = dat.get_grid() - values = dat.get_values() - num_dims = len(np.array(values.shape).squeeze()) - - growth_rates = np.zeros(1) - ks = np.zeros(1) - if num_dims == 2: - if kwargs["dir"] == 0: - growth_rates = np.zeros(values.shape[1]) - ks = np.zeros(values.shape[1]) - elif kwargs["dir"] == 1: - growth_rates = np.zeros(values.shape[0]) - ks = np.zeros(values.shape[0]) - # end - # end - - for idx in range(len(growth_rates)): - p0 = kwargs["guess"] - if kwargs["guess"]: - guess = kwargs["guess"].split(",") - p0 = (float(guess[0]), float(guess[1])) - # end - - x = time[0] - if kwargs["dir"] == 1: - x = time[1] - - y = values[..., 0].squeeze() - if kwargs["dir"] == 0: - y = values[:, idx, 0].squeeze() - elif kwargs["dir"] == 1: - y = values[idx, :, 0].squeeze() - # end - - best_params, _, _ = postgkyl.tools.fit_growth(x, y, min_N=kwargs["minn"], p0=p0) - - if kwargs["dataset"]: - out = GData(tag="growth", label="Fit", - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - t = 0.5 * (time[0][:-1] + time[0][1:]) - out_val = postgkyl.tools.exp2(t, *best_params) - out.push([time[0]], out_val[..., np.newaxis]) - data.add(out) - # end - - if kwargs["instantaneous"]: - verb_print(ctx, "growth: Plotting instantaneous growth rate") - gammas = [] - for i in range(1, len(time[0]) - 1): - gamma = (values[i + 1, 0] - values[i - 1, 0]) / (2*values[i, 0]*(time[0][i + 1] - time[0][i - 1])) - gammas.append(gamma) - - plt.style.use(f"{os.path.dirname(os.path.realpath(__file__)):s}/../output/postgkyl.mplstyle") - _, ax = plt.subplots() - ax.plot(time[0][1:-1], gammas) - # ax.set_autoscale_on(False) - ax.grid(True) - plt.show() - # end - # end - - growth_rates[idx] = best_params[1] - ks[idx] = idx - # end - - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - out.push([ks], growth_rates[..., np.newaxis]) - data.add(out) - # end - # end - verb_print(ctx, "Finishing growth") diff --git a/src/postgkyl/commands/info.py b/src/postgkyl/commands/info.py deleted file mode 100644 index ab62e273..00000000 --- a/src/postgkyl/commands/info.py +++ /dev/null @@ -1,37 +0,0 @@ -import click - -from postgkyl.utils import verb_print - - -@click.command(help="Print info of active datasets.") -@click.option("-u", "--use", help="Specify a 'tag' to apply to (default all tags).") -@click.option("-c", "--compact", is_flag=True, help="Show in compact mode.") -@click.option("-a", "--allsets", is_flag=True, help="All data sets.") -@click.pass_context -def info(ctx, **kwargs): - verb_print(ctx, "Starting info") - data = ctx.obj["data"] - if kwargs["allsets"]: - only_active = False - else: - only_active = True - # end - - for i, dat in data.iterator(kwargs["use"], enum=True, only_active=only_active): - if dat.get_status(): - color = "green" - bold = True - else: - color = None - bold = False - # end - click.echo( - click.style(f"{dat.get_label():s}{' ' if dat.get_label() else '':s}({dat.get_tag():s}#{i:d})", - fg=color, bold=bold) - ) - if not kwargs["compact"]: - click.echo(dat.info() + "\n") - # end - # end - - verb_print(ctx, "Finishing info") diff --git a/src/postgkyl/commands/integrate.py b/src/postgkyl/commands/integrate.py deleted file mode 100644 index b14a784a..00000000 --- a/src/postgkyl/commands/integrate.py +++ /dev/null @@ -1,32 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -import postgkyl.tools as tools - - -@click.command() -@click.argument("axis", nargs=1, type=click.STRING) -@click.option("--use", "-u", default=None, help="Specify the tag to integrate.") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def integrate(ctx, **kwargs): - """"Integrate data over a specified axis or axes.""" - verb_print(ctx, "Starting integrate") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - if kwargs["tag"]: - grid, values = tools.integrate(dat, kwargs["axis"]) - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - out.push(grid, values) - data.add(out) - else: - tools.integrate(dat, kwargs["axis"], overwrite=True) - # end - # end - - verb_print(ctx, "Finishing integrate") diff --git a/src/postgkyl/commands/interpolate.py b/src/postgkyl/commands/interpolate.py deleted file mode 100644 index d6e38ee5..00000000 --- a/src/postgkyl/commands/interpolate.py +++ /dev/null @@ -1,81 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.data import GInterpModal, GInterpNodal -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--basis_type","-b", - type=click.Choice(["ms", "ns", "mo", "mt", "gkhyb", "gkhyb_vel", "pkpmhyb"]), - help="Specify DG basis.") -@click.option("--poly_order", "-p", type=click.INT, help="Specify polynomial order.") -@click.option("--interp", "-i", type=click.INT, - help="Interpolation onto a general mesh of specified amount.") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.option("--read", "-r", type=click.BOOL, help="Read from general interpolation file.") -@click.pass_context -def interpolate(ctx, **kwargs): - """Interpolate DG data onto a uniform mesh.""" - verb_print(ctx, "Starting interpolate") - data = ctx.obj["data"] - - basis_type = None - is_modal = None - if kwargs.get("basis_type"): - if kwargs["basis_type"] == "ms": - basis_type = "serendipity" - is_modal = True - elif kwargs["basis_type"] == "ns": - basis_type = "serendipity" - is_modal = False - elif kwargs["basis_type"] == "mo": - basis_type = "maximal-order" - is_modal = True - elif kwargs["basis_type"] == "mt": - basis_type = "tensor" - is_modal = True - elif kwargs["basis_type"] == "gkhyb": - basis_type = "gkhybrid" - is_modal = True - elif kwargs["basis_type"] == "gkhyb_vel": - basis_type = "gkhybrid_vel" - is_modal = True - elif kwargs["basis_type"] == "pkpmhyb": - basis_type = "hybrid" - is_modal = True - # end - # end - - for dat in data.iterator(kwargs["use"]): - if kwargs["basis_type"] is None and dat.ctx["basis_type"] is None: - ctx.fail( - click.style(f"ERROR in interpolate: no 'basis_type' was specified and dataset {dat.get_label():s} does not have required ctxdata", - fg="red") - ) - # end - - - if is_modal or dat.ctx["is_modal"]: - dg = GInterpModal(dat, kwargs["poly_order"], kwargs["basis_type"], - kwargs["interp"], kwargs["read"]) - else: - dg = GInterpNodal(dat, kwargs["poly_order"], basis_type, kwargs["interp"], kwargs["read"]) - # end - - num_nodes = dg.num_nodes - num_comps = int(dat.get_num_comps() / num_nodes) - - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - grid, values = dg.interpolate(tuple(range(num_comps))) - out.push(grid, values) - data.add(out) - else: - dg.interpolate(tuple(range(num_comps)), overwrite=True) - # end - # end - verb_print(ctx, "Finishing interpolate") diff --git a/src/postgkyl/commands/laguerre_compose.py b/src/postgkyl/commands/laguerre_compose.py deleted file mode 100644 index 23405b8e..00000000 --- a/src/postgkyl/commands/laguerre_compose.py +++ /dev/null @@ -1,29 +0,0 @@ -import click - -from postgkyl.data import GData -import postgkyl.tools -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--distribution", "-f", type=click.STRING, prompt=True, - help="Specify the PKPM distribution function dataset.") -@click.option("--tm", type=click.STRING, prompt=True, help="Specify the PKPM vars dataset.") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.pass_context -def laguerrecompose(ctx, **kwargs): - """Compose PKPM Laguerre coefficients together.""" - verb_print(ctx, "Starting laguerrecompose") - data = ctx.obj["data"] - - for f, tm in zip(data.iterator(kwargs["distribution"]), data.iterator(kwargs["tm"])): - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=f.ctx) - postgkyl.tools.laguerre_compose(f, tm, out) - else: - postgkyl.tools.laguerre_compose(f, tm, f) - # end - # end - verb_print(ctx, "Finishing laguerrecompose") diff --git a/src/postgkyl/commands/listoutputs.py b/src/postgkyl/commands/listoutputs.py deleted file mode 100644 index e49603eb..00000000 --- a/src/postgkyl/commands/listoutputs.py +++ /dev/null @@ -1,43 +0,0 @@ -from glob import glob -import click -import re - -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--extensions", "-e", type=click.STRING, default="bp,gkyl", - show_default=True, help="Output file extension(s)") -@click.option("--path", "-p", type=click.Path(exists=True, file_okay=False), - default=".", show_default=True, help="Path to search for outputs") -@click.pass_context -def listoutputs(ctx, **kwargs): - """List Gkeyll filename stems in the current directory.""" - verb_print(ctx, "Starting listoutputs") - - extensions = kwargs["extensions"].split(",") - path = kwargs["path"] - for ext in extensions: - files = glob(f"{path}/*.{ext:s}") - unique = [] - for fn in files: - # remove extension - s = fn[: -(len(ext) + 1)] - # strip "restart" - if s.endswith("_restart"): - s = s[:-8] - # end - # strip digits - s = re.sub(r"_\d+$", "", s) - if s not in unique: - unique.append(s) - # end - # end - if len(unique) > 0: - click.echo(f"{ext:s}:") - # end - for s in sorted(unique): - click.echo(f"- {s:s}") - # end - # end - verb_print(ctx, "Finishing listoutputs") diff --git a/src/postgkyl/commands/load.py b/src/postgkyl/commands/load.py deleted file mode 100644 index 40d61631..00000000 --- a/src/postgkyl/commands/load.py +++ /dev/null @@ -1,155 +0,0 @@ -import click -import glob - -from postgkyl.data import GData -from postgkyl.data import GInterpModal -from postgkyl.utils import verb_print - - -def _pick_cut(ctx : click.Context, kwargs : dict, zn : int) -> str | None: - nm = f"z{zn:d}" - if zn == 6: # This little hack allows to apply the same function for - # components as well - nm = "component" - # end - if kwargs[nm] and ctx.obj["global_cuts"][zn]: - click.echo(click.style(f"WARNING: The local '{nm:s}' is overwriting the global '{nm:s}'", - fg="yellow")) - return kwargs[nm] - elif kwargs[nm]: - return kwargs[nm] - elif ctx.obj["global_cuts"][zn]: - return ctx.obj["global_cuts"][zn] - else: - return None - # end - - -def _crush(s : str) -> tuple: # Temp function used as a sorting key - splitted = s.split("_") - tmp = splitted[-1].split(".") - splitted[-1] = int(tmp[0]) - splitted.append(tmp[1]) - return tuple(splitted) - - -@click.command(hidden=True) -@click.option("--z0", help="Partial file load: 0th coord (either int or slice).") -@click.option("--z1", help="Partial file load: 1st coord (either int or slice).") -@click.option("--z2", help="Partial file load: 2nd coord (either int or slice).") -@click.option("--z3", help="Partial file load: 3rd coord (either int or slice).") -@click.option("--z4", help="Partial file load: 4th coord (either int or slice).") -@click.option("--z5", help="Partial file load: 5th coord (either int or slice).") -@click.option("--component", "-c", help="Partial file load: comps (either int or slice).") -@click.option("--tag", "-t", default="default", help="Specily tag for data.") -@click.option("--compgrid", is_flag=True, help="Disregard the mapped grid information") -@click.option("--varname", "-d", multiple=True, - help="Allows to specify the Adios variable name. [default: 'CartGridField']") -@click.option("--label", "-l", help="Allows to specify the custom label") -@click.option("--c2p", type=click.STRING, - help="Specify the file name containing c2p mapped coordinates") -@click.option("--c2p-vel", "c2p_vel",type=click.STRING, - help="Specify the file name containing c2p mapped coordinates") -@click.option("--fv", is_flag=True, - help="Tag finite volume data when using c2p mapped coordinates") -@click.option("--reader", "-r", type=click.STRING, - help="Allows to specify the Adios variable name (default is 'CartGridField')") -@click.option("--load/--no-load", default=True, help="Specify if data should be loaded.") -@click.pass_context -def load(ctx, **kwargs): - verb_print(ctx, "Starting load") - data = ctx.obj["data"] - - idx = ctx.obj["in_data_strings_loaded"] - in_data_string = ctx.obj["in_data_strings"][idx] - - # Handling the wildcard characters - if "*" in in_data_string or "?" in in_data_string or "!" in in_data_string: - files = glob.glob(str(in_data_string)) - files = [f for f in files if f.find("restart") < 0] - try: - files = sorted(files, key=_crush) - except Exception: - click.echo( - click.style("WARNING: The loaded files appear to be of different types. Sorting is turned off.", - fg="yellow") - ) - # end - else: - files = [in_data_string] - # end - - # Resolve the local/global variable names and partial loading - # The local settings take a precedents but a warning is going to appear - z0 = _pick_cut(ctx, kwargs, 0) - z1 = _pick_cut(ctx, kwargs, 1) - z2 = _pick_cut(ctx, kwargs, 2) - z3 = _pick_cut(ctx, kwargs, 3) - z4 = _pick_cut(ctx, kwargs, 4) - z5 = _pick_cut(ctx, kwargs, 5) - comp = _pick_cut(ctx, kwargs, 6) - - var_names = ["CartGridField"] - if kwargs["varname"] and ctx.obj["global_var_names"]: - var_names = kwargs["varname"] - click.echo( - click.style("WARNING: The local 'varname' is overwriting the global 'varname'", - fg="yellow") - ) - elif kwargs["varname"]: - var_names = kwargs["varname"] - elif ctx.obj["global_var_names"]: - var_names = ctx.obj["global_var_names"] - # end - - mapc2p_name = None - if kwargs["c2p"] and ctx.obj["global_c2p"]: - mapc2p_name = kwargs["c2p"] - click.echo( - click.style("WARNING: The local 'c2p' is overwriting the global 'c2p'", fg="yellow") - ) - elif kwargs["c2p"]: - mapc2p_name = kwargs["c2p"] - elif ctx.obj["global_c2p"]: - mapc2p_name = ctx.obj["global_c2p"] - # end - - mapc2p_vel_name = None - if kwargs["c2p_vel"] and ctx.obj["global_c2p_vel"]: - mapc2p_name = kwargs["c2p_vel"] - click.echo( - click.style("WARNING: The local 'c2p_vel' is overwriting the global 'c2p_vel'", - fg="yellow") - ) - elif kwargs["c2p_vel"]: - mapc2p_vel_name = kwargs["c2p_vel"] - elif ctx.obj["global_c2p_vel"]: - mapc2p_vel_name = ctx.obj["global_c2p_vel"] - # end - - if len(var_names) == 1: - var_names = var_names[0].split(",") - # end - - for var in var_names: - for fn in files: - try: - dat = GData(file_name=fn, tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - z0=z0, z1=z1, z2=z2, z3=z3, z4=z4, z5=z5, comp=comp, var_name=var, - label=kwargs["label"], mapc2p_name=mapc2p_name, mapc2p_vel_name=mapc2p_vel_name, - reader_name=kwargs["reader"], load=kwargs["load"], click_mode=True) - if kwargs["fv"]: - dg = GInterpModal(dat, 0, "ms") - dg.interpolateGrid(overwrite=True) - # end - data.add(dat) - except NameError as e: - ctx.fail(click.style(rf"{repr(e):s}", fg="red")) - # end - # end - # end - - data.set_unique_labels() - - ctx.obj["in_data_strings_loaded"] += 1 - verb_print(ctx, "Finishing load") diff --git a/src/postgkyl/commands/magsq.py b/src/postgkyl/commands/magsq.py deleted file mode 100644 index 40c179c1..00000000 --- a/src/postgkyl/commands/magsq.py +++ /dev/null @@ -1,30 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -import postgkyl.tools - - -@click.command() -@click.option("--use", "-u", default=None, help="Specify the tag to integrate.") -@click.option("--tag", "-t", default=None, help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def magsq(ctx, **kwargs): - """Calculate the magnitude squared of an input array.""" - verb_print(ctx, "Starting magnitude squared computation") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - postgkyl.tools.mag_sq(dat, output=out) - data.add(out) - else: - postgkyl.tools.mag_sq(dat, output=dat) - # end - # end - - verb_print(ctx, "Finishing magnitude squared computation") diff --git a/src/postgkyl/commands/mask.py b/src/postgkyl/commands/mask.py deleted file mode 100644 index 795f5432..00000000 --- a/src/postgkyl/commands/mask.py +++ /dev/null @@ -1,45 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--filename", "-f", type=click.STRING, help="Specify the file with a mask.") -@click.option("--lower", "-l", type=click.FLOAT, - help="Specify the lower theshold to be masked out.") -@click.option("--upper", "-u", type=click.FLOAT, - help="Specify the upper theshold to be masked out.") -@click.pass_context -def mask(ctx, **kwargs): - """Mask data with specified Gkeyll mask file.""" - verb_print(ctx, "Starting mask") - data = ctx.obj("data") - - if kwargs["filename"]: - mask_fld = GData(kwargs["filename"]).get_values() - # end - - for dat in data.interator(kwargs["use"]): - values = dat.get_values() - - if kwargs["filename"]: - mask_fld_rep = np.repeat(mask_fld, dat.get_num_comps(), axis=-1) - data.set_values(np.ma.masked_where(mask_fld_rep < 0.0, values)) - elif kwargs.get("lower") and kwargs.get("upper"): - dat.set_values(np.ma.masked_outside(values, kwargs["lower"], kwargs["upper"])) - elif kwargs.get("lower"): - dat.set_values(np.ma.masked_less(values, kwargs["lower"])) - elif kwargs.get("upper"): - dat.set_values(np.ma.masked_greater(values, kwargs["upper"])) - else: - data.set_values(values) - click.echo( - click.style("WARNING in 'mask': No masking information specified.", fg="yellow") - ) - # end - # end - - verb_print(ctx, "Finishing mask") diff --git a/src/postgkyl/commands/mhd.py b/src/postgkyl/commands/mhd.py deleted file mode 100644 index 78fca8b1..00000000 --- a/src/postgkyl/commands/mhd.py +++ /dev/null @@ -1,67 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -import postgkyl.tools.prim_vars as pv - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--mu0", "-m", type=click.FLOAT, default=1.0, show_default=True, - help="Permeability of free space.") -@click.option("--gas_gamma", "-g", type=click.FLOAT, default=5.0/3, show_default=True, - help="Gas adiabatic constant.") -@click.option("--variable_name", "-v", prompt=True, - type=click.Choice(["density", "xvel", "yvel", "zvel", "vel", "Bx", "By", "Bz", "Bi", - "magpressure", "pressure", "temp", "sound", "mach"]), - help="Variable to extract") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.pass_context -def mhd(ctx, **kwargs): - """Compute ideal MHD primitive and some derived variables from MHD conserved variables. - """ - verb_print(ctx, "Starting mhd") - data = ctx.obj["data"] - - v = kwargs["variable_name"] - for dat in data.iterator(kwargs["use"]): - verb_print(ctx, f"mhd: Extracting {v:s} from data set") - out = dat - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - data.add(out) - # end - if v == "density": - pv.get_density(dat, out_mom=out) - elif v == "xvel": - pv.get_vx(dat, out_mom=out) - elif v == "yvel": - pv.get_vy(dat, out_mom=out) - elif v == "zvel": - pv.get_vz(dat, out_mom=out) - elif v == "vel": - pv.get_vi(dat, out_mom=out) - elif v == "Bx": - pv.get_mhd_Bx(dat, out_mom=out) - elif v == "By": - pv.get_mhd_By(dat, out_mom=out) - elif v == "Bz": - pv.get_mhd_Bz(dat, out_mom=out) - elif v == "Bi": - pv.get_mhd_Bi(dat, out_mom=out) - elif v == "magpressure": - pv.get_mhd_mag_p(dat, mu_0=kwargs["mu0"], out_mom=out) - elif v == "pressure": - pv.get_mhd_p(dat, gas_gamma=kwargs["gas_gamma"], mu_0=kwargs["mu0"], out_mom=out) - elif v == "temp": - pv.get_mhd_temp(dat, gas_gamma=kwargs["gas_gamma"], mu_0=kwargs["mu0"], out_mom=out) - elif v == "sound": - pv.get_mhd_sound(dat, gas_gamma=kwargs["gas_gamma"], mu_0=kwargs["mu0"], out_mom=out) - elif v == "mach": - pv.get_mhd_mach(dat, gas_gamma=kwargs["gas_gamma"], mu_0=kwargs["mu0"], out_mom=out) - # end - # end - verb_print(ctx, "Finishing mhd") diff --git a/src/postgkyl/commands/old/cglpressure.py b/src/postgkyl/commands/old/cglpressure.py deleted file mode 100644 index 0b0fd8be..00000000 --- a/src/postgkyl/commands/old/cglpressure.py +++ /dev/null @@ -1,103 +0,0 @@ -import click -import numpy as np - -from postgkyl.commands import tm -from postgkyl.tools.stack import pushStack, peakStack, antiSqueeze, addStack -from postgkyl.utils import verb_print - - - -def getParPerp(pij, B): - tmp = np.copy(pij[..., 0:2]) - - pxx = pij[..., 0] - pxy = pij[..., 1] - pxz = pij[..., 2] - pyy = pij[..., 3] - pyz = pij[..., 4] - pzz = pij[..., 5] - - b = np.sqrt(B[..., 0] * B[..., 0] + B[..., 1] * B[..., 1] + B[..., 2] * B[..., 2]) - bx = B[..., 0] / b - by = B[..., 1] / b - bz = B[..., 2] / b - - tmp[..., 0] = ( - bx * bx * pxx - + by * by * pyy - + bz * bz * pzz - + 2.0 * (bx * by * pxy + bx * bz * pxz + by * bz * pyz) - ) - tmp[..., 1] = (pxx + pyy + pzz - tmp[..., 0]) / 2.0 - - return tmp - - -def getAgyro(pij, B): - tmp = np.copy(pij[..., 0:6]) - - pxx = pij[..., 0] - pxy = pij[..., 1] - pxz = pij[..., 2] - pyy = pij[..., 3] - pyz = pij[..., 4] - pzz = pij[..., 5] - - b = np.sqrt(B[..., 0] * B[..., 0] + B[..., 1] * B[..., 1] + B[..., 2] * B[..., 2]) - bx = B[..., 0] / b - by = B[..., 1] / b - bz = B[..., 2] / b - - ppar = ( - bx * bx * pxx - + by * by * pyy - + bz * bz * pzz - + 2.0 * (bx * by * pxy + bx * bz * pxz + by * bz * pyz) - ) - pper = (pxx + pyy + pzz - ppar) / 2.0 - - tmp[..., 0] = pxx - (ppar * bx * bx + pper * (1 - bx * bx)) # xx - tmp[..., 1] = pxy - (ppar * bx * by + pper * (0 - bx * by)) # xy - tmp[..., 2] = pxz - (ppar * bx * bz + pper * (0 - bx * bz)) # xz - tmp[..., 3] = pyy - (ppar * by * by + pper * (1 - by * by)) # yy - tmp[..., 4] = pyz - (ppar * by * bz + pper * (0 - by * bz)) # yz - tmp[..., 5] = pzz - (ppar * bz * bz + pper * (1 - bz * bz)) # zz - - return tmp - - -@click.command() -@click.option( - "--agyro", - is_flag=True, - default=False, - help="Compute the agyrotropic part of pressure tensor instead", -) -@click.pass_context -def cglpressure(ctx, **inputs): - """Extract parallel and perpendicular pressures from pressure-tensor - and magnetic field. Pressure-tensor must be the first dataset and - magnetic field the second dataset. A two component field - (parallel, perpendicular) is returned. Optionally, the command can - extract the six components of the agyrotropic part of the pressure - tensor. - - """ - verb_print(ctx, "Starting CGL pressure") - - coords, pij = peakStack(ctx, ctx.obj["sets"][0]) - coords, B = peakStack(ctx, ctx.obj["sets"][1]) - - if inputs["agyro"]: - tmp = getAgyro(pij, B) - else: - tmp = getParPerp(pij, B) - - tmp = antiSqueeze(coords, tmp) - - idx = addStack(ctx) - ctx.obj["type"].append("hist") - pushStack(ctx, idx, coords, tmp, "CGL") - ctx.obj["sets"] = [idx] - - verb_print(ctx, "Finishing CGL pressure") diff --git a/src/postgkyl/commands/old/recovery.py b/src/postgkyl/commands/old/recovery.py deleted file mode 100644 index 8f436c9a..00000000 --- a/src/postgkyl/commands/old/recovery.py +++ /dev/null @@ -1,64 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GInterpModal -from postgkyl.utils import verb_print - -from postgkyl.data import GData - - -@click.command(help="Interpolate DG data on a uniform mesh") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.option( - "--basis_type", "-b", type=click.Choice(["ms", "ns", "mo"]), help="Specify DG basis" -) -@click.option("--poly_order", "-p", type=click.INT, help="Specify polynomial order") -@click.option("--interp", "-i", type=click.INT, help="Number of poins to evaluate on") -@click.option( - "-r", "--periodic", is_flag=True, help="Flag for periodic boundary conditions" -) -@click.option("-c", "--c1", is_flag=True, help="Enforce continuous first derivatives") -@click.pass_context -def recovery(ctx, **kwargs): - verb_print(ctx, "Starting recovery") - data = ctx.obj["data"] - - if "basis_type" in kwargs.keys(): - if kwargs["basis_type"] == "ms" or kwargs["basis_type"] == "ns": - basis_type = "serendipity" - elif kwargs["basis_type"] == "mo": - basis_type = "maximal-order" - # end - else: - basis_type = None - # end - - for dat in data.iterator(kwargs["use"]): - dg = GInterpModal( - dat, kwargs["poly_order"], basis_type, kwargs["interp"], kwargs["periodic"] - ) - num_nodes = dg.num_nodes - num_comps = int(dat.get_num_comps() / num_nodes) - - # verb_print(ctx, 'interplolate: interpolating dataset #{:d}'.format(s)) - # dg.recovery(tuple(range(num_comps)), stack=True) - if kwargs["tag"]: - out = GData( - tag=kwargs["tag"], - label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], - ctx=dat.ctx, - ) - grid, values = dg.recovery(0, kwargs["c1"]) - out.push(grid, values) - data.add(out) - else: - dg.recovery(0, kwargs["c1"], overwrite=True) - # end - # end - verb_print(ctx, "Finishing recovery") - - -# end diff --git a/src/postgkyl/commands/parrotate.py b/src/postgkyl/commands/parrotate.py deleted file mode 100644 index d0af42dc..00000000 --- a/src/postgkyl/commands/parrotate.py +++ /dev/null @@ -1,42 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.parrotate - - -@click.command() -@click.option("--array", "-a", default="array", show_default=True, - help="Tag for array to be rotated") -@click.option("--rotator", "-r", default="rotator", show_default=True, - help="Tag for rotator (data used for the rotation)") -@click.option("--tag", "-t", default="rotarraypar", show_default=True, - help="Tag for the resulting rotated array parallel to rotator") -@click.option("--label", "-l", default="rotarraypar", show_default=True, - help="Custom label for the result") -@click.pass_context -def parrotate(ctx, **kwargs): - """Rotate an array parallel to the unit vectors of a second array. - - For two arrays u and v, where v is the rotator, operation is (u dot v_hat) v_hat. Note - that for a three-component field, the output is a new vector whose components are - (u_{v_x}, u_{v_y}, u_{v_z}), i.e., the x, y, and z components of the vector u parallel - to v. - """ - verb_print(ctx, "Starting rotation parallel to rotator array") - - data = ctx.obj["data"] - - for a, rot in zip(data.iterator(kwargs["array"]), data.iterator(kwargs["rotator"])): - grid, outrot = postgkyl.tools.parrotate(a, rot) - # Create new GData structure with appropriate outtag and labels to store output. - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=a.ctx) - out.push(outrot, grid) - data.add(out) - # end - - data.deactivate_all(tag=kwargs["array"]) - data.deactivate_all(tag=kwargs["rotator"]) - - verb_print(ctx, "Finishing rotation parallel to rotator array") diff --git a/src/postgkyl/commands/perprotate.py b/src/postgkyl/commands/perprotate.py deleted file mode 100644 index 90e8b3f1..00000000 --- a/src/postgkyl/commands/perprotate.py +++ /dev/null @@ -1,39 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.perprotate - - -@click.command() -@click.option("--array", "-a", default="array", show_default=True, - help="Tag for array to be rotated") -@click.option("--rotator", "-r", default="rotator", show_default=True, - help="Tag for rotator (data used for the rotation)") -@click.option("--tag", "-t", default="rotarrayperp", show_default=True, - help="Tag for the resulting rotated array perpendicular to rotator") -@click.option("--label", "-l", default="rotarrayperp", show_default=True, - help="Custom label for the result") -@click.pass_context -def perprotate(ctx, **kwargs): - """Rotate an array perpendicular to the unit vectors of a second array. - - For two arrays u and v, where v is the rotator, operation is u - (u dot v_hat) v_hat. - """ - verb_print(ctx, "Starting rotation perpendicular to rotator array") - - data = ctx.obj["data"] # shortcut - - for a, rot in zip(data.iterator(kwargs["array"]), data.iterator(kwargs["rotator"])): - grid, outrot = postgkyl.tools.perprotate(a, rot) - # Create new GData structure with appropriate outtag and labels to store output. - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=a.ctx) - out.push(outrot, grid) - data.add(out) - # end - - data.deactivate_all(tag=kwargs["array"]) - data.deactivate_all(tag=kwargs["rotator"]) - - verb_print(ctx, "Finishing rotation perpendicular to rotator array") diff --git a/src/postgkyl/commands/plot.py b/src/postgkyl/commands/plot.py deleted file mode 100644 index 41b50d2e..00000000 --- a/src/postgkyl/commands/plot.py +++ /dev/null @@ -1,343 +0,0 @@ -import click -import matplotlib.pyplot as plt -import numpy as np - -from postgkyl.utils import verb_print -import postgkyl.output.plot - - -@click.command() -@click.option("--use", "-u", default=None, help="Specify the tag to plot.") -@click.option("--figure", "-f", default=None, - help="Specify figure to plot in; either number or 'dataset'.") -@click.option("--squeeze", is_flag=True, help="Squeeze the components into one panel.") -@click.option("--subplots", "-b", is_flag=True, help="Make subplots from multiple datasets.") -@click.option("--nsubplotrow", "num_subplot_row", type=click.INT, - help="Manually set the number of rows for subplots.") -@click.option("--nsubplotcol", "num_subplot_col", type=click.INT, - help="Manually set the number of columns for subplots.") -@click.option("--transpose", is_flag=True, help="Transpose axes.") -@click.option("-c", "--contour", is_flag=True, help="Make contour plot.") -@click.option("--surface", "--surf", "surface", is_flag=True, - help="Make a 3D surface plot for 2D data (auto-enabled when overlaying " - "multiple 2D datasets).") -@click.option("--alpha", type=click.FLOAT, default=None, - help="Surface transparency (0-1); useful when overlaying surfaces.") -@click.option("--clevels", type=click.STRING, - help="Specify levels for contours: comma-separated level values or start:end:nlevels.") -@click.option("--cnlevels", type=click.INT, help="Specify the number of levels for contours.") -@click.option("--contlabel", "cont_label", is_flag=True, help="Add labels to contours") -@click.option("-q", "--quiver", is_flag=True, help="Make quiver plot.") -@click.option("-l", "--streamline", is_flag=True, help="Make streamline plot.") -@click.option("--sdensity", type=click.INT, default=1, help="Control density of the streamlines.") -@click.option("--arrowstyle", type=click.STRING, help="Set the style for streamline arrows.") -@click.option("--lineouts", type=click.Choice(["0", "1"]), help="Switch to lineouts mode.") -@click.option("-s", "--scatter", is_flag=True, help="Make scatter plot.") -@click.option("--markersize", type=click.FLOAT, help="Set marker size for scatter plots.") -@click.option("--linewidth", type=click.FLOAT, help="Set the linewidth.") -@click.option("--linestyle", type=click.Choice(["solid", "dashed", "dotted", "dashdot"]), - help="Set the linestyle.") -@click.option("--style", help="Specify Matplotlib style file (default: Postgkyl).") -@click.option("-d", "--diverging", is_flag=True, help="Switch to diverging color map.") -@click.option("--arg", type=click.STRING, default="", - help="Additional plotting arguments, e.g., '*--'.") -@click.option("--fix-aspect", "-a", "fixaspect", is_flag=True, - help="Enforce the same scaling on both axes.") -@click.option("--aspect", default=None, help="Specify the scaling ratio.") -@click.option("--logx", is_flag=True, help="Set x-axis to log scale.") -@click.option("--logy", is_flag=True, help="Set y-axis to log scale.") -@click.option("--logz", is_flag=True, help="Set values of 2D plot to log scale.") -@click.option("--xshift", default=0.0, type=click.FLOAT, show_default=True, - help="Value to shift the x-axis.") -@click.option("--yshift", default=0.0, type=click.FLOAT, show_default=True, - help="Value to shift the y-axis.") -@click.option("--zshift", default=0.0, type=click.FLOAT, show_default=True, - help="Value to shift the z-axis.") -@click.option("--xscale", default=1.0, type=click.FLOAT, show_default=True, - help="Value to scale the x-axis.") -@click.option("--yscale", default=1.0, type=click.FLOAT, show_default=True, - help="Value to scale the y-axis.") -@click.option("--zscale", default=1.0, type=click.FLOAT, show_default=True, - help="Value to scale the z-axis (default: 1.0).") -@click.option("--xmax", default=None, type=click.FLOAT, help="Set maximal x-value.") -@click.option("--xmin", default=None, type=click.FLOAT, help="Set minimal x-values.") -@click.option("--ymax", default=None, type=click.FLOAT, help="Set maximal y-value.") -@click.option("--ymin", default=None, type=click.FLOAT, help="Set minimal y-values.") -@click.option("--zmax", default=None, type=click.FLOAT, help="Set maximal z-value.") -@click.option("--zmin", default=None, type=click.FLOAT, help="Set minimal z-values.") -@click.option("--xlim", default=None, type=click.STRING, - help="Set limits for the x-coordinate (lower,upper)") -@click.option("--ylim", default=None, type=click.STRING, - help="Set limits for the y-coordinate (lower,upper).") -@click.option("--zlim", default=None, type=click.STRING, - help="Set limits for the z-coordinate (lower,upper).") -@click.option("--relax", is_flag=True, help="Relax the stringent x axis limits for 1D plots.") -@click.option("--globalrange", "-r", is_flag=True, help="Make uniform extends across datasets.") -@click.option("--cutoffglobalrange", "-cogr", default=None, type=click.FLOAT, - help="Set custom limit for uniform across datasets") -@click.option("--legend", default=None, type=click.STRING, - help="If specified, comma-separated legend labels (e.g., 'a,b,c').") -@click.option("--no-legend", is_flag=True, help="Hide legend.") -@click.option("--force-legend", "forcelegend", is_flag=True, - help="Force legend even when plotting a single dataset.") -@click.option("--color", type=click.STRING, help="Set color when available.") -@click.option("-x", "--xlabel", type=click.STRING, help="Specify a x-axis label.") -@click.option("-y", "--ylabel", type=click.STRING, help="Specify a y-axis label.") -@click.option("--clabel", type=click.STRING, help="Specify a label for colorbar.") -@click.option("--title", type=click.STRING, help="Specify a title.") -@click.option("--subplot-titles", type=click.STRING, help="Comma-separated titles for each subplot. e.g. --subplot-titles 'Title1,Title2,Title3'") -@click.option("--subplot-xlabels", type=click.STRING, help="Comma-separated x-axis labels for each subplot. e.g. --subplot-xlabels 'X1,X2,X3'") -@click.option("--subplot-ylabels", type=click.STRING, help="Comma-separated y-axis labels for each subplot. e.g. --subplot-ylabels 'Y1,Y2,Y3'") -@click.option("--save", is_flag=True, help="Save figure as PNG file.") -@click.option("--saveas", type=click.STRING, default=None, help="Name of figure file.") -@click.option("--dpi", type=click.INT, default=200, help="DPI (resolution) for output.") -@click.option("-e", "--edgecolors", type=click.STRING, - help="Set color for cell edges to show grid outline.") -@click.option("--showgrid/--no-showgrid", default=True, help="Show grid-lines.") -@click.option("--xkcd", is_flag=True, help="Turns on the xkcd style!") -@click.option("--hashtag", is_flag=True, help="Turns on the pgkyl hashtag!") -@click.option("--show/--no-show", default=True, - help="Turn showing of the plot ON and OFF.") -@click.option("--figsize", help="Comma-separated values for x and y size.") -@click.option("--saveframes", type=click.STRING, - help="Save individual frames as PNGS instead of an opening them") -@click.option("--jet", is_flag=True, help="Turn colormap to jet for comparison with literature.") -@click.option("--cmap", "--colormap", type=click.STRING, default=None, - help="Override default colormap with a valid matplotlib cmap.") -@click.option("--cval", type=click.STRING, default=None, - help="For 1D plots, comma-separated values mapping each curve onto the colormap " - "(e.g. '1e-6,2e-6'). Requires --cmap; defaults to the dataset index if omitted.") -@click.option("-m", "--multiblock", is_flag=True, default=False) -@click.pass_context -def plot(ctx, **kwargs): - """Plot active datasets, optionally displaying the plot and/or saving it to PNG files. - - Plot labels can use a sub-set of LaTeX math commands placed between dollar ($) signs. - """ - verb_print(ctx, "Starting plot") - - kwargs["rcParams"] = ctx.obj["rcParams"] - - args = kwargs["arg"] - if kwargs["scatter"]: - args += "." - # end - del kwargs["arg"] - - if kwargs["jet"]: - click.echo( - click.style("WARNING: The 'jet' colormap has been selected. This colormap is not perceptually uniform and seemingly creates features which do not exist in the data!", - fg="yellow") - ) - # end - - if kwargs["aspect"]: - kwargs["fixaspect"] = True - # end - - if kwargs["lineouts"]: - kwargs["lineouts"] = int(kwargs["lineouts"]) - # end - - kwargs["num_axes"] = None - if kwargs["subplots"]: - kwargs["num_axes"] = 0 - kwargs["start_axes"] = 0 - for dat in ctx.obj["data"].iterator(kwargs["use"]): - kwargs["num_axes"] = kwargs["num_axes"] + dat.get_num_comps() - # end - if kwargs["figure"] is None: - kwargs["figure"] = 0 - # end - # end - - if kwargs["xlim"]: - kwargs["xmin"] = float(kwargs["xlim"].split(",")[0]) - kwargs["xmax"] = float(kwargs["xlim"].split(",")[1]) - # end - if kwargs["ylim"]: - kwargs["ymin"] = float(kwargs["ylim"].split(",")[0]) - kwargs["ymax"] = float(kwargs["ylim"].split(",")[1]) - # end - if kwargs["zlim"]: - kwargs["zmin"] = float(kwargs["zlim"].split(",")[0]) - kwargs["zmax"] = float(kwargs["zlim"].split(",")[1]) - # end - - dataset_fignum = False - if ( - kwargs["figure"] == "dataset" - or kwargs["figure"] == "set" - or kwargs["figure"] == "s" - ): - dataset_fignum = True - # end - - #automatically sets correct scale for multiblock cases - if kwargs["multiblock"] and kwargs["cutoffglobalrange"] is None: - kwargs["globalrange"] = True - # end - - # When several 2D datasets are drawn into the same figure we switch to contour mode. - num_datasets = sum(1 for _ in ctx.obj["data"].iterator(kwargs["use"])) - first_dat = next(ctx.obj["data"].iterator(kwargs["use"]), None) - is_2d = first_dat is not None and first_dat.get_num_dims(squeeze=True) == 2 - overlay_2d = ( - is_2d and num_datasets > 1 and not dataset_fignum - and kwargs["figure"] is not None - and not kwargs["subplots"] and kwargs["lineouts"] is None - and not kwargs["quiver"] and not kwargs["streamline"] - ) - if overlay_2d and not kwargs["surface"] and not kwargs["contour"]: - kwargs["contour"] = True - # end - kwargs["comparison"] = overlay_2d and (kwargs["surface"] or kwargs["contour"]) - - if kwargs["globalrange"] or kwargs["cutoffglobalrange"]: - vmin = float("inf") - vmax = float("-inf") - v_extrema = np.array([]) - for dat in ctx.obj["data"].iterator(kwargs["use"]): - val = dat.get_values() * kwargs["zscale"] - if vmin > np.nanmin(val): - vmin = np.nanmin(val) - # end - if vmax < np.nanmax(val): - vmax = np.nanmax(val) - # end - v_extrema = np.append(v_extrema, np.nanmin(val)) - v_extrema = np.append(v_extrema, np.nanmax(val)) - # end - - v_extrema = np.sort(v_extrema) - if kwargs["cutoffglobalrange"]: - boundary = 100 * (1 - kwargs["cutoffglobalrange"]) / 2 - vmax = np.percentile(v_extrema, 100 - boundary) - vmin = np.percentile(v_extrema, boundary) - # end - - if kwargs["zmin"] is None: - kwargs["zmin"] = vmin - # end - if kwargs["zmax"] is None: - kwargs["zmax"] = vmax - # end - # end - - #Prevents scale errors for multiblock contour plots - if kwargs["multiblock"] and kwargs["contour"] and kwargs["clevels"] is None: - kwargs["clevels"] = f"{kwargs['zmin']}:{kwargs['zmax']}:10" - # end - - # Parse legend labels if provided - legend_labels = None - if kwargs.get("legend"): - legend_labels = [label.strip() for label in kwargs["legend"].split(",")] - # end - - # Overwrite show_legend if no_legend is set - show_legend = True - if kwargs.get("no_legend"): - if kwargs["no_legend"]: - show_legend = False - - kwargs["legend"] = show_legend - del kwargs["no_legend"] - - # Colormap based line coloring for 1D plots. - cval_list = None - if kwargs["cval"]: - cval_list = [float(v) for v in kwargs["cval"].split(",")] - elif kwargs["cmap"]: - num_datasets = sum(1 for _ in ctx.obj["data"].iterator(kwargs["use"])) - cval_list = list(range(num_datasets)) - # end - del kwargs["cval"] - if cval_list: - kwargs["cval_min"] = min(cval_list) - kwargs["cval_max"] = max(cval_list) - else: - kwargs["cval_min"] = None - kwargs["cval_max"] = None - # end - - file_name = "" - - # ---- Loop over all the datasets ---- - for i, dat in ctx.obj["data"].iterator(kwargs["use"], enum=True): - if dataset_fignum: - kwargs["figure"] = int(i) - # end - #puts all blocks on the same figure - if kwargs["multiblock"]: - kwargs["figure"] = 0 - # end - - # Determine the label for this dataset - if legend_labels is not None and i < len(legend_labels): - label = legend_labels[i] - elif ctx.obj["data"].get_num_datasets() > 1 or kwargs["forcelegend"]: - label = dat.get_label() - else: - label = "" - # end - - # 1D colouring. - if cval_list is not None and i < len(cval_list): - kwargs["cval"] = cval_list[i] - else: - kwargs["cval"] = None - # end - - # ---- Plot ---- - postgkyl.output.plot(dat, args, label_prefix=label, **kwargs) - - if kwargs["subplots"]: - kwargs["start_axes"] = kwargs["start_axes"] + dat.get_num_comps() - # end - - if kwargs["save"] or kwargs["saveas"]: - if kwargs["saveas"]: - file_name = kwargs["saveas"] - else: - if file_name != "": - file_name = file_name + "_" - # end - if dat._file_name: - file_name = file_name + dat._file_name.split(".")[0] - else: - file_name = file_name + "ev_" + ctx.obj["labels"][i].replace(" ", "_") - # end - # end - # end - if (kwargs["save"] or kwargs["saveas"]) and kwargs["figure"] is None: - file_name = str(file_name) - plt.savefig(file_name, dpi=kwargs["dpi"]) - file_name = "" - # end - - if kwargs["saveframes"]: - file_name = f"{kwargs['saveframes']:s}_{i:d}.png" - plt.savefig(file_name, dpi=kwargs["dpi"]) - kwargs["show"] = False - # end - - if "batch_mode" in ctx.obj: - if ctx.obj["batch_mode"]: - file_name = f"{ctx.obj['saveframes_prefix']:s}_{i:d}.png" - plt.savefig(file_name, dpi=kwargs["dpi"]) - kwargs["show"] = False - # end - # end - - - # end - if (kwargs["save"] or kwargs["saveas"]): - file_name = str(file_name) - plt.savefig(file_name, dpi=kwargs["dpi"]) - # end - - if kwargs["show"]: - plt.show() - # end - verb_print(ctx, "Finishing plot") diff --git a/src/postgkyl/commands/pr.py b/src/postgkyl/commands/pr.py deleted file mode 100644 index 50ae8380..00000000 --- a/src/postgkyl/commands/pr.py +++ /dev/null @@ -1,27 +0,0 @@ -import click -import numpy as np - -from postgkyl.utils import verb_print - -np.set_printoptions(precision=16) - -@click.command(help="Print the data") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--grid", "-g", is_flag=True, help="Print grid instead of values.") -@click.pass_context -def pr(ctx, **kwargs): - verb_print(ctx, "Starting pr") - data = ctx.obj["data"] - - for dat in data.iterator(kwargs["use"]): - if kwargs["grid"]: - grid = dat.get_grid() - for g in grid: - click.echo(g) - # end - else: - click.echo(dat.get_values().squeeze()) - # end - # end - - verb_print(ctx, "Finishing pr") diff --git a/src/postgkyl/commands/relchange.py b/src/postgkyl/commands/relchange.py deleted file mode 100644 index f838547a..00000000 --- a/src/postgkyl/commands/relchange.py +++ /dev/null @@ -1,36 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.tools.rel_change - - -@click.command(help="Computes the relative change between two datasets") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--index", "-i", type=click.INT, default=0, show_default=True, - help="Dataset index for computing change relative to.") -@click.option("--comp", "-c", default=None, show_default=True, - help="Dataset component to be compared to if user only wants to compare to a single component.") -@click.option("--tag", "-t", default="rel_change", show_default=True, help="Tag for the result.") -@click.option("--label", "-l", default="delta", show_default=True, help="Custom label for the result/") -@click.pass_context -def relchange(ctx, **kwargs): - verb_print(ctx, "Starting relative change") - - data = ctx.obj["data"] - for tag in data.tag_iterator(kwargs["use"]): - reference = data.get_dataset(kwargs["index"], tag) - for dat in data.iterator(tag): - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], compgrid=ctx.obj["compgrid"], ctx=dat.ctx) - grid, values = postgkyl.tools.rel_change(reference, dat, kwargs["comp"]) - dat.deactivate() - out.push(grid, values) - data.add(out) - else: - grid, values = postgkyl.tools.rel_change(reference, dat, kwargs["comp"]) - dat.push(grid, values) - # end - # end - # end - verb_print(ctx, "Finishing relative change") diff --git a/src/postgkyl/commands/select.py b/src/postgkyl/commands/select.py deleted file mode 100644 index 5c636ecf..00000000 --- a/src/postgkyl/commands/select.py +++ /dev/null @@ -1,162 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print, set_frame - -import postgkyl.data.select - - -@click.command() -@click.option("--z0", default=None, help="Indices for 0th coord (either int, float, or slice).") -@click.option("--z1", default=None, help="Indices for 1st coord (either int, float, or slice).") -@click.option("--z2", default=None, help="Indices for 2nd coord (either int, float, or slice).") -@click.option("--z3", default=None, help="Indices for 3rd coord (either int, float, or slice).") -@click.option("--z4", default=None, help="Indices for 4th coord (either int, float, or slice).") -@click.option("--z5", default=None, help="Indices for 5th coord (either int, float, or slice).") -@click.option("--comp", "-c", default=None, - help="Indices for components (either int, slice, or coma-separated).") -@click.option("--use", "-u", help="Specify a 'tag' to apply to.") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result") -@click.option("--multiblock", "-m", is_flag=True, - help="Necessary parameter for multiblock lineouts in z0 or z1 dims") -@click.option("--multiframe", "-f", is_flag=True, - help="Specify if performing select on multiple multiblock frames") -@click.pass_context -def select(ctx, **kwargs): - """Subselect data from the active dataset(s). - - This command allows, for example, to choose a specific component of a multi-component - dataset, select a index or coordinate range. Index ranges can also be specified using - python slice notation (start:end:stride). - """ - verb_print(ctx, "Starting select") - data = ctx.obj["data"] - - #multiblock case - if kwargs["multiblock"]: - - #set ctx frames - frame_list = set_frame(ctx) - #creates list of lists with blocks per frame if multiframe parameter - #if not, then only one frame with all blocks - if kwargs["multiframe"]: - data_list = [] - for frame in frame_list: - frame_data_list = [dat for dat in data.iterator(kwargs["use"]) if dat.ctx["frame"] == frame] - data_list.append(frame_data_list) - # end - else: - data_list = [list(data.iterator(kwargs["use"]))] - # end - - - for i, frame in enumerate(data_list): - - #establish lower bounds for x and y axis - botlef_point = [] - for dim in [0,1]: - botlef_point.append(min([dat.get_bounds()[0][dim] for dat in frame])) - # end - #find starting block for lineout coordinate - if kwargs.get("z0"): - for dat in frame: - if dat.get_bounds()[0][0] <= float(kwargs["z0"]) <= dat.get_bounds()[1][0] and dat.get_bounds()[0][1] == botlef_point[1]: - block = dat - # end - # end - # end - if kwargs.get("z1"): - for dat in frame: - if dat.get_bounds()[0][1] <= float(kwargs["z1"]) <= dat.get_bounds()[1][1] and dat.get_bounds()[0][0] == botlef_point[0]: - block = dat - # end - # end - # end - #find neighboring blocks of starting block - block.set_neighbors(frame) - - value_list = [] - - #creates new grid and value list containing data from blocks which contain specified z0 coordinate - if kwargs.get("z0"): - grid, values = postgkyl.data.select(block, - z0=kwargs["z0"], - comp=kwargs["comp"]) - grid_list = grid - for val in values[0]: - value_list.append(val) - # end - while block._neighbors[1][1] is not None: - block = block._neighbors[1][1] - block.set_neighbors(data.iterator(kwargs["use"])) - grid, values = postgkyl.data.select(block, - z0=kwargs["z0"], - comp=kwargs["comp"]) - grid_list[1] = np.append(grid_list[1], grid[1]) - for val in values[0]: - value_list.append(val) - # end - # end - grid_list[1] = np.unique(grid_list[1]) - value_list = np.array([value_list]) - # end - - - #same but for z1 coordinate - if kwargs.get("z1"): - grid, values = postgkyl.data.select(block, - z1=kwargs["z1"], - comp=kwargs["comp"]) - grid_list = grid - for val in values: - value_list.append(val) - # end - while block._neighbors[0][1] is not None: - block = block._neighbors[0][1] - block.set_neighbors(data.iterator(kwargs["use"])) - grid, values = postgkyl.data.select(block, - z1=kwargs["z1"], - comp=kwargs["comp"]) - grid_list[0] = np.append(grid_list[0], grid[0]) - for val in values: - value_list.append(val) - # end - grid_list[0] = np.unique(grid_list[0]) - value_list = np.array(value_list) - # end - - #loop through frame list and deactivate each - for dat in frame: - dat.deactivate() - # end - - #create new gdata instance and push new stitched grid and values - out = GData(tag=kwargs["tag"], - label=kwargs["label"], - comp_grid=ctx.obj["compgrid"]) - out.ctx["frame"] = i - out.push(grid_list, value_list) - data.add(out) - # end - - - else: - for dat in data.iterator(kwargs["use"]): - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - grid, values = postgkyl.data.select(dat, - z0=kwargs["z0"], z1=kwargs["z1"], z2=kwargs["z2"], z3=kwargs["z3"], - z4=kwargs["z4"], z5=kwargs["z5"], comp=kwargs["comp"]) - out.push(grid, values) - data.add(out) - else: - postgkyl.data.select(dat, overwrite=True, - z0=kwargs["z0"], z1=kwargs["z1"], z2=kwargs["z2"], z3=kwargs["z3"], - z4=kwargs["z4"], z5=kwargs["z5"], comp=kwargs["comp"]) - # end - # end - # end - verb_print(ctx, "Finishing select") diff --git a/src/postgkyl/commands/status.py b/src/postgkyl/commands/status.py deleted file mode 100644 index 776a7764..00000000 --- a/src/postgkyl/commands/status.py +++ /dev/null @@ -1,71 +0,0 @@ -import click - -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--tag", "-t", type=click.STRING, help="Tag(s) to apply to (comma-separated).") -@click.option("--index", "-i", type=click.STRING, - help="Dataset indices (e.g., '1', '0,2,5', or '1:6:2').") -@click.option("--focused", "-f", is_flag=True, help="Leave unspecified datasets untouched.") -@click.pass_context -def activate(ctx, **kwargs): - """Select datasets(s) to pass further down the command chain. - - Datasets are indexed starting 0. Multiple datasets can be selected using a comma - separated list or a range specifier. Unless '--focused' is selected, all unselected - datasets will be deactivated. - - '--tag' and '--index' allow to specify tags and indices. The not specified, 'activate' - applies to all. Both parameters support comma-separated values. '--index' also - supports slices following the Python conventions, e.g., '3:7' or ':-5:2'. - - 'info' command (especially with the '-ac' flags) can be helpful when - activating/deactivating multiple datasets. - """ - verb_print(ctx, "Starting activate") - data = ctx.obj["data"] - - if not kwargs["focused"]: - data.deactivate_all() - # end - - for dat in data.iterator(tag=kwargs["tag"], only_active=False, select=kwargs["index"]): - dat.activate() - # end - - verb_print(ctx, "Finishing activate") - - -@click.command() -@click.option("--tag", "-t", type=click.STRING, help="Tag(s) to apply to (comma-separated).") -@click.option("--index", "-i", type=click.STRING, - help="Dataset indices (e.g., '1', '0,2,5', or '1:6:2').") -@click.option("--focused", "-f", is_flag=True, help="Leave unspecified datasets untouched.") -@click.pass_context -def deactivate(ctx, **kwargs): - """Select datasets(s) to pass further down the command chain. - - Datasets are indexed starting 0. Multiple datasets can be selected using a comma - separated list or a range specifier. Unless '--focused' is selected, all unselected - datasets will be activated. - - '--tag' and '--index' allow to specify tags and indices. The not specified, - 'deactivate' applies to all. Both parameters support comma-separated values. '--index' - also supports slices following the Python conventions, e.g., '3:7' or ':-5:2'. - - 'info' command (especially with the '-ac' flags) can be helpful when - activating/deactivating multiple datasets. - """ - verb_print(ctx, "Starting deactivate") - data = ctx.obj["data"] - - if kwargs["focused"]: - data.activate_all() - # end - - for dat in data.iterator(tag=kwargs["tag"], only_active=False, select=kwargs["index"]): - dat.deactivate() - # end - - verb_print(ctx, "Finishing deactivate") diff --git a/src/postgkyl/commands/style.py b/src/postgkyl/commands/style.py deleted file mode 100644 index 4bcfec74..00000000 --- a/src/postgkyl/commands/style.py +++ /dev/null @@ -1,35 +0,0 @@ -import click - -from postgkyl.utils import load_style, verb_print - - -@click.command() -@click.option("--file", "-f", help="Sets Maplotlib rcParams style file.") -@click.option("--set", "-s", multiple=True, help="Sets individual rcParam(s) as 'key:value'.") -@click.option("--print", "-p", is_flag=True, help="Prints the current rcParams.") -@click.pass_context -def style(ctx, **kwargs): - """Probe and control the Matplotlib plotting style. - - The list of rcParams is available - here:\nhttps://matplotlib.org/stable/api/matplotlib_configuration_api.html""" - verb_print(ctx, "Starting 'style' command") - - if kwargs["file"]: - load_style(ctx, kwargs["file"]) - # end - - for param in kwargs["set"]: - param_split = param.split(":") - key = param_split[0].strip() - value = param[len(param_split[0]) + 1 :].strip() - ctx.obj["rcParams"][key] = value - # end - - if kwargs["print"]: - for key in ctx.obj["rcParams"]: - print(f"{key:s} : {ctx.obj['rcParams'][key]}") - # end - # end - - verb_print(ctx, "Finishing 'style' command") diff --git a/src/postgkyl/commands/temp.py b/src/postgkyl/commands/temp.py deleted file mode 100644 index 1d4db9f0..00000000 --- a/src/postgkyl/commands/temp.py +++ /dev/null @@ -1,76 +0,0 @@ -import click -import numpy as np - -from postgkyl.utils import verb_print - - - -# ---- Math ---- -@click.command(help="Multiply data by a factor") -@click.argument("factor", nargs=1, type=click.FLOAT) -@click.pass_context -def mult(ctx, **kwargs): - verb_print(ctx, f"Multiplying by {kwargs['factor']:f}") - for s in ctx.obj["sets"]: - values = ctx.obj["dataSets"][s].get_values() - values = values * kwargs["factor"] - ctx.obj["dataSets"][s].push(values) - # end - - -@click.command(help="Calculate power of data") -@click.argument("power", nargs=1, type=click.FLOAT) -@click.pass_context -def pow(ctx, **kwargs): - verb_print(ctx, f"Calculating the power of {kwargs['power']:f}") - for s in ctx.obj["sets"]: - values = ctx.obj["dataSets"][s].get_values() - values = values ** kwargs["power"] - ctx.obj["dataSets"][s].push(values) - # end - - -@click.command(help="Calculate natural log of data") -@click.pass_context -def log(ctx): - verb_print(ctx, "Calculating the natural log") - for s in ctx.obj["sets"]: - values = ctx.obj["dataSets"][s].get_values() - values = np.log(values) - ctx.obj["dataSets"][s].push(values) - # end - - -@click.command(help="Calculate absolute values of data") -@click.pass_context -def abs(ctx): - verb_print(ctx, "Calculating the absolute value") - for s in ctx.obj["sets"]: - values = ctx.obj["dataSets"][s].get_values() - values = np.abs(values) - ctx.obj["dataSets"][s].push(values) - # end - - -@click.command(help="Normalize data") -@click.option("--shift/--no-shift", default=False, show_default=True, - help="Shift minimal value to zero.") -@click.option("--usefirst", is_flag=True, default=False, help="Normalize to first value in field.") -@click.pass_context -def norm(ctx, **kwargs): - verb_print(ctx, "Normalizing data") - for s in ctx.obj["sets"]: - values = ctx.obj["dataSets"][s].get_values() - num_comps = ctx.obj["dataSets"][s].get_num_comps() - values_out = values.copy() - for comp in range(num_comps): - if kwargs["shift"]: - values_out[..., comp] -= values_out[..., comp].min() - if kwargs["usefirst"]: - values_out[..., comp] /= values_out[..., comp].item(0) - else: - values_out[..., comp] /= np.abs(values_out[..., comp]).max() - # end - # end - ctx.obj["dataSets"][s].push(values_out) - # end diff --git a/src/postgkyl/commands/tenmoment.py b/src/postgkyl/commands/tenmoment.py deleted file mode 100644 index e5324abf..00000000 --- a/src/postgkyl/commands/tenmoment.py +++ /dev/null @@ -1,71 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print - -import postgkyl.tools.prim_vars as pv - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("-v", "--variable_name", prompt=True, - type=click.Choice(["density", "xvel", "yvel", "zvel", "vel", "pressureTensor", - "pxx", "pxy", "pxz", "pyy", "pyz", "pzz", "pressure", "temp", "ke", "sound", "mach"]), - help="Variable to work with.") -@click.option("-g", "--gas_gamma",type=click.FLOAT, show_default=True, default=5.0/3, - help="Gas adiabatic constant.") -@click.option("--tag", "-t", help="Optional tag for the resulting array") -@click.option("--label", "-l", help="Custom label for the result") -@click.pass_context -def tenmoment(ctx, **kwargs): - """Extract ten-moment primitive variables from ten-moment conserved variables. - """ - verb_print(ctx, "Starting tenmoment") - data = ctx.obj["data"] - - v = kwargs["variable_name"] - for dat in data.iterator(kwargs["use"]): - verb_print(ctx, f"tenmoment: Extracting {v:s} from data set") - out = dat - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - data.add(out) - # end - if v == "density": - pv.get_density(dat, out_mom=out) - elif v == "xvel": - pv.get_vx(dat, out_mom=out) - elif v == "yvel": - pv.get_vy(dat, out_mom=out) - elif v == "zvel": - pv.get_vz(dat, out_mom=out) - elif v == "vel": - pv.get_vi(dat, out_mom=out) - elif v == "pressureTensor": - pv.get_pij(dat, out_mom=out) - elif v == "pxx": - pv.get_pxx(dat, out_mom=out) - elif v == "pxy": - pv.get_pxy(dat, out_mom=out) - elif v == "pxz": - pv.get_pxz(dat, out_mom=out) - elif v == "pyy": - pv.get_pyy(dat, out_mom=out) - elif v == "pyz": - pv.get_pyz(dat, out_mom=out) - elif v == "pzz": - pv.get_pzz(dat, out_mom=out) - elif v == "pressure": - pv.get_p(dat, gas_gamma=kwargs["gas_gamma"], num_moms=10, out_mom=out) - elif v == "ke": - pv.get_ke(dat, gas_gamma=kwargs["gas_gamma"], num_moms=10, out_mom=out) - elif v == "temp": - pv.get_temp(dat, gas_gamma=kwargs["gas_gamma"], num_moms=10, out_mom=out) - elif v == "sound": - pv.get_sound(dat, gas_gamma=kwargs["gas_gamma"], num_moms=10, out_mom=out) - elif v == "mach": - pv.get_mach(dat, gas_gamma=kwargs["gas_gamma"], num_moms=10, out_mom=out) - # end - # end - verb_print(ctx, "Finishing tenmoment") diff --git a/src/postgkyl/commands/trajectory.py b/src/postgkyl/commands/trajectory.py deleted file mode 100644 index 9125ed5e..00000000 --- a/src/postgkyl/commands/trajectory.py +++ /dev/null @@ -1,138 +0,0 @@ -from matplotlib.animation import FuncAnimation -import click -import math -import matplotlib.pyplot as plt -import numpy as np - -from postgkyl.utils import verb_print - - - -def _update(i, ax, ctx, leap, vel, xmin, xmax, ymin, ymax, zmin, zmax, tag): - colors = ["C0", "C1", "C2", "C3", "C4", "C5", "C6", "C7", "C8", "C9"] - - s = 0 - plt.cla() - # for s, dat in ctx.obj['data'].iterator(tag, emum=True): - for dat in ctx.obj["data"].iterator(tag): - time = dat.get_grid()[0] - coords = dat.get_values() - t_idx = int(i * leap) - - if xmin is not None: - x = np.where(coords[:, 0] > xmin, coords[:, 0], np.nan) - else: - x = coords[:, 0] - # end - if xmax is not None: - x = np.where(x < xmax, x, np.nan) - # end - if ymin is not None: - y = np.where(coords[:, 1] > ymin, coords[:, 1], np.nan) - else: - y = coords[:, 1] - # end - if ymax is not None: - y = np.where(y < ymax, y, np.nan) - # end - if zmin is not None: - z = np.where(coords[:, 2] > zmin, coords[:, 2], np.nan) - else: - z = coords[:, 2] - # end - if zmax is not None: - z = np.where(z < zmax, z, np.nan) - # end - - ax.plot(x, y, z, color=colors[s % 10]) - ax.scatter(x[t_idx], y[t_idx], z[t_idx], color=colors[s % 10]) - if vel and dat.get_num_comps() == 6: - if t_idx + leap >= len(time): - dt = time[-1] - time[t_idx] - else: - dt = time[int(t_idx + leap)] - time[t_idx] - # end - dx = coords[i, 3] * dt - dy = coords[i, 4] * dt - dz = coords[i, 5] * dt - ax.plot([x[t_idx], x[t_idx] + dx], [y[t_idx], y[t_idx] + dy], [z[t_idx], z[t_idx] + dz], - color=colors[s % 10]) - # end - s += 1 - # end - plt.title(f"T: {time[t_idx]:.4e}") - ax.set_xlabel("$z_0$") - ax.set_ylabel("$z_1$") - ax.set_zlabel("$z_2$") - ax.set_xlim3d(xmin, xmax) - ax.set_ylim3d(ymin, ymax) - ax.set_zlim3d(zmin, zmax) - - -@click.command() -@click.option("--fix-aspect", "fixaspect",is_flag=True, help="Enforce the same scaling on both axes.") -@click.option("--show/--no-show", default=True, help="Turn showing of the plot ON and OFF (default: ON).") -@click.option("-i", "--interval", default=100, help="Specify the animation interval.") -@click.option("--save", is_flag=True, help="Save figure as PNG.") -@click.option("--velocity/--no-velocity", default=True, help="Plot velocity vectors.") -@click.option("--saveas", type=click.STRING, default=None, help="Name to save the plot as.") -@click.option("-e", "--elevation", type=click.FLOAT, help="Set elevation.") -@click.option("-a", "--azimuth", type=click.FLOAT, help="Set azimuth.") -@click.option("-n", "--numframes", type=click.INT, help="Set number of frames for the animation.") -@click.option("--xmin", type=click.FLOAT, help="Minimum value of the x-coordinate") -@click.option("--xmax", type=click.FLOAT, help="Maximum value of the x-coordinate") -@click.option("--ymin", type=click.FLOAT, help="Minimum value of the y-coordinate") -@click.option("--ymax", type=click.FLOAT, help="Maximum value of the y-coordinate") -@click.option("--zmin", type=click.FLOAT, help="Minimum value of the z-coordinate") -@click.option("--zmax", type=click.FLOAT, help="Maximum value of the z-coordinate") -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.pass_context -def trajectory(ctx, **kwargs): - """Animate a particle trajectory.""" - verb_print(ctx, "Starting trajectory") - data = ctx.obj["data"] - - tags = list(data.tag_iterator(kwargs["use"])) - tag = tags[0] - if len(tags) > 1: - ctx.fail(click.echo(f"'trajectory' supports only one 'tag', was provided {len(tags):d}", - color="red")) - # end - - fig = plt.figure() - ax = fig.add_subplot(111, projection="3d") - kwargs["figure"] = fig - kwargs["legend"] = False - - dat = ctx.obj["data"].get_dataset(0, tag) - num_pos = dat.get_num_cells()[0] - - jump = 1 - if kwargs.get("numframes"): - jump = int(math.floor(num_pos / kwargs["numframes"])) - num_pos = int(kwargs["numframes"]) - # end - - anim = FuncAnimation(fig, _update, num_pos, - fargs=(ax, ctx, jump, kwargs["velocity"], kwargs["xmin"], kwargs["xmax"], kwargs["ymin"], - kwargs["ymax"], kwargs["zmin"], kwargs["zmax"], tag), - interval=kwargs["interval"]) - - ax.view_init(elev=kwargs["elevation"], azim=kwargs["azimuth"]) - - if kwargs["fixaspect"]: - plt.setp(ax, aspect=1.0) - # end - - f_name = "anim.mp4" - if kwargs["saveas"]: - f_name = str(kwargs["saveas"]) - # end - if kwargs["save"] or kwargs["saveas"]: - anim.save(f_name, writer="ffmpeg") - # end - - if kwargs["show"]: - plt.show() - # end - verb_print(ctx, "Finishing trajectory") diff --git a/src/postgkyl/commands/transform_frame.py b/src/postgkyl/commands/transform_frame.py deleted file mode 100644 index ec6fa1a0..00000000 --- a/src/postgkyl/commands/transform_frame.py +++ /dev/null @@ -1,32 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.tools import transform_frame -from postgkyl.utils import verb_print - - - -@click.command() -@click.option("--distribution", "-f", type=click.STRING, prompt=True, - help="Specify the PKPM distribution function.") -@click.option("--bulk", "-u", type=click.STRING, prompt=True, help="Specify the PKPM moments.") -@click.option("--cdim", "-c", type=click.INT, prompt=True, - help="Specify the number of configuration space dimensions.") -@click.option("--tag", "-t", help="Optional tag for the resulting array.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.pass_context -def transformframe(ctx, **kwargs): - """Compose PKPM Laguerre coefficients together.""" - verb_print(ctx, "Starting transformframe") - data = ctx.obj["data"] - - for f, bulk in zip(data.iterator(kwargs["distribution"]), data.iterator(kwargs["bulk"])): - if kwargs["tag"]: - out = GData(tag=kwargs["tag"], label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=f.ctx) - transform_frame(f, bulk, kwargs["cdim"], out) - else: - transform_frame(f, bulk, kwargs["cdim"], f) - # end - # end - verb_print(ctx, "Finishing transformframe") diff --git a/src/postgkyl/commands/val2coord.py b/src/postgkyl/commands/val2coord.py deleted file mode 100644 index fdcdac00..00000000 --- a/src/postgkyl/commands/val2coord.py +++ /dev/null @@ -1,110 +0,0 @@ -import click -import numpy as np - -from postgkyl.data import GData -from postgkyl.utils import verb_print - - -def _get_range(str_in, length): - if len(str_in.split(",")) > 1: - return np.array(str_in.split(","), np.int) - elif str_in.find(":") >= 0: - str_split = str_in.split(":") - - if str_split[0] == "": - s_idx = 0 - else: - s_idx = int(str_split[0]) - if s_idx < 0: - s_idx = length + s_idx - # end - # end - - if str_split[1] == "": - e_idx = length - else: - e_idx = int(str_split[1]) - if e_idx < 0: - e_idx = length + e_idx - # end - # end - - inc = 1 - if len(str_split) > 2 and str_split[2] != "": - inc = int(str_split[2]) - # end - return np.arange(s_idx, e_idx, inc) - else: - return np.array([int(str_in)]) - # end - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("--tag", "-t", help="Tag for the result.") -@click.option("--label", "-l", help="Custom label for the result.") -@click.option("-x", type=click.STRING, - help="Select components that will became the grid of the new dataset.") -@click.option("-y", type=click.STRING, - help="Select components that will became the values of the new dataset.") -@click.option("--periodic", "-p", is_flag=True, help="Set the last component to match the first one.") -@click.pass_context -def val2coord(ctx, **kwargs): - """Given a dataset (typically a DynVector) selects columns from it to create new datasets. - - For example, you can choose say column 1 to be the X-axis of the new dataset and - column 2 to be the Y-axis. Multiple columns can be choosen using range specifiers and - as many datasets are then created. - """ - verb_print(ctx, "Starting val2coord") - data = ctx.obj["data"] - - tags = list(data.tag_iterator()) - out_tag = kwargs["tag"] - if out_tag is None: - if len(tags) == 1: - out_tag = tags[0] - else: - out_tag = "val2coord" - # end - # end - - for _, dat in data.iterator(kwargs["use"], enum=True): - values = dat.get_values() - x_comps = _get_range(kwargs["x"], len(values[0, :])) - y_comps = _get_range(kwargs["y"], len(values[0, :])) - - if len(x_comps) > 1 and len(x_comps) != len(y_comps): - click.echo( - click.style(f"ERROR 'val2coord': Length of the x-components ({len(x_comps):d}) is greater than 1 and not equal to the y-components ({len(y_comps):d}).", - fg="red") - ) - ctx.exit() - # end - - for i, yc in enumerate(y_comps): - if len(x_comps) > 1: - xc = x_comps[i] - else: - xc = x_comps[0] - # end - - x = values[..., xc] - y = values[..., yc] - - if kwargs["periodic"]: - x = np.append(x, np.atleast_1d(x[0]), axis=0) - y = np.append(y, np.atleast_1d(y[0]), axis=0) - # end - - y = y[..., np.newaxis] # Adding the required component index - - out = GData(tag=out_tag, label=kwargs["label"], - comp_grid=ctx.obj["compgrid"], ctx=dat.ctx) - out.push([x], y) - out.color = "C0" - data.add(out) - # end - dat.deactivate() - # end - verb_print(ctx, "Finishing val2coord") diff --git a/src/postgkyl/commands/velocity.py b/src/postgkyl/commands/velocity.py deleted file mode 100644 index 14e18e38..00000000 --- a/src/postgkyl/commands/velocity.py +++ /dev/null @@ -1,33 +0,0 @@ -import click - -from postgkyl.data import GData -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--density", "-d", default="density", show_default=True, help="Tag for density.") -@click.option("--momentum", "-m", default="momentum", show_default=True, help="Tag for momentum.") -@click.option("--tag", "-t", default="velocity", show_default=True, help="Tag for the result.") -@click.option("--label", "-l", default="velocity", show_default=True, - help="Custom label for the result.") -@click.pass_context -def velocity(ctx, **kwargs): - verb_print(ctx, "Starting velocity") - - data = ctx.obj["data"] # shortcut - - for m0, m1 in zip(data.iterator(kwargs["density"]), data.iterator(kwargs["momentum"])): - grid = m0.get_grid() - vals_M0 = m0.get_values() - vals_M1 = m1.get_values() - - out = GData(tag=kwargs["tag"], comp_grid=ctx.obj["compgrid"], - label=kwargs["label"], ctx=m0.ctx) - out.push(grid, vals_M1 / vals_M0) - data.add(out) - # end - - data.deactivate_all(tag=kwargs["density"]) - data.deactivate_all(tag=kwargs["momentum"]) - - verb_print(ctx, "Finishing velocity") diff --git a/src/postgkyl/commands/write.py b/src/postgkyl/commands/write.py deleted file mode 100644 index 271b34e6..00000000 --- a/src/postgkyl/commands/write.py +++ /dev/null @@ -1,58 +0,0 @@ -import click -import shutil - -from postgkyl.utils import verb_print - - -@click.command() -@click.option("--use", "-u", help="Specify a 'tag' to apply to (default all tags).") -@click.option("-f", "--filename", type=click.STRING, prompt=True, help="Output file name.") -@click.option("-m", "--mode", type=click.Choice(["gkyl", "bp", "txt", "npy"]), default="gkyl", - help="Output file mode. One of `gkyl` (binary, default), `bp` (ADIOS BP file), `txt` (ASCII text file), or `npy` (NumPy binary file).") -@click.option("-s", "--single", is_flag=True, help="Write all dataset into one file") -@click.pass_context -def write(ctx, **kwargs): - """Write active dataset to a file. - - The output file format can be set with ``--format``, and is Gkeyll's .gkyl by default. - Files saved as .gkyl or .bp can be later loaded back into pgkyl to further manipulate - or plot. - """ - verb_print(ctx, "Starting write") - data = ctx.obj["data"] - - var_name = None - append = False - cleaning = True - fn = kwargs["filename"] - mode = kwargs["mode"] - if len(fn.split(".")) > 1: - mode = str(fn.split(".")[-1]) - fn = str(fn.split(".")[0]) - # end - - num_files = data.get_num_datasets(tag=kwargs["use"]) - for i, dat in data.iterator(tag=kwargs["use"], enum=True): - out_name = f"{fn:s}.{mode:s}" - if kwargs["single"]: - var_name = f"{dat.get_tag():s}_{i:d}" - cleaning = False - else: - if num_files > 1: - out_name = f"{fn:s}_{i:d}.{mode:s}" - # end - # end - - dat.write(out_name=out_name, mode=mode, append=append, var_name=var_name, cleaning=cleaning) - - if kwargs["single"]: - append = True - # end - # end - - # Cleaning - if not cleaning: - shutil.move(f"{fn:s}.{mode:s}.dir/{fn:s}.{mode:s}.0", f"{fn:s}.{mode:s}") - shutil.rmtree(f"{fn:s}.{mode:s}.dir") - # end - verb_print(ctx, "Finishing write") diff --git a/src/postgkyl/data/__init__.py b/src/postgkyl/data/__init__.py deleted file mode 100644 index 081fad9e..00000000 --- a/src/postgkyl/data/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# Import data handler -from .gdata import GData - -# Import interpolators -from .dg import GInterpNodal -from .dg import GInterpModal - -# Import interpolation matrices computation -from . import computeInterpolationMatrices -from . import computeDerivativeMatrices - -# Import select -from .select import select - -from .idx_parser import idx_parser - -from .gkyl_reader import GkylReader -from .gkyl_adios_reader import GkylAdiosReader -from .gkyl_h5_reader import GkylH5Reader -from .flash_h5_reader import FlashH5Reader diff --git a/src/postgkyl/data/computeDerivativeMatrices.py b/src/postgkyl/data/computeDerivativeMatrices.py deleted file mode 100644 index c1c57e13..00000000 --- a/src/postgkyl/data/computeDerivativeMatrices.py +++ /dev/null @@ -1,7948 +0,0 @@ -import numpy -from sympy import * - -from optparse import OptionParser - - -def createDerivativeMatrix(dim, order, basis_type, interp, modal=True): - interpFloat = float(interp) - interpList = numpy.zeros(interp) - - for i in range(0, interpList.shape[0]): - interpList[i] = ( - -1.0 * (interpFloat - 1) / interpFloat + float(i) * 2.0 / interpFloat - ) - - if dim == 1: - x = Symbol("x") - if modal: - if order == 1: - - functionVector = Matrix([[0.7071067811865468], [1.224744871391589 * x]]) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - - elif order == 2: - - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - elif order == 3: - - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - [4.677071733467427 * x**3 - 2.806243040080457 * x], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - - elif order == 4: - - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - [4.677071733467427 * x**3 - 2.806243040080457 * x], - [ - 9.280776503073431 * x**4 - - 7.954951288348656 * x**2 - + 0.7954951288348655 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - - functionVector = Matrix([[0.5 - 0.5 * x], [0.5 + 0.5 * x]]) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - - elif order == 2: - - functionVector = Matrix( - [[0.5 * x**2 - 0.5 * x], [1.0 - x**2], [0.5 * x**2 + 0.5 * x]] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - elif order == 3: - - functionVector = Matrix( - [ - [-(9.0 * x**3) / 16.0 + (9.0 * x**2) / 16.0 + x / 16.0 - 1 / 16.0], - [ - (27.0 * x**3) / 16.0 - - (9.0 * x**2) / 16.0 - - (27.0 * x) / 16.0 - + 9.0 / 16.0 - ], - [ - (27.0 * x) / 16.0 - - (9.0 * x**2) / 16.0 - - (27.0 * x**3) / 16.0 - + 9.0 / 16.0 - ], - [(9.0 * x**3) / 16.0 + (9.0 * x**2) / 16.0 - x / 16.0 - 1 / 16.0], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - - elif order == 4: - - functionVector = Matrix( - [ - [(2.0 * x**4) / 3.0 - (2.0 * x**3) / 3.0 - x**2 / 6.0 + x / 6.0], - [ - -(8.0 * x**4) / 3.0 - + (4.0 * x**3) / 3.0 - + (8.0 * x**2) / 3.0 - - (4.0 * x) / 3.0 - ], - [4.0 * x**4 - 5.0 * x**2 + 1.0], - [ - -(8.0 * x**4) / 3.0 - - (4.0 * x**3) / 3.0 - + (8.0 * x**2) / 3.0 - + (4.0 * x) / 3.0 - ], - [(2.0 * x**4) / 3.0 + (2.0 * x**3) / 3.0 - x**2 / 6.0 - x / 6.0], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i] = diff(functionVector[i], x) - - derivativeMatrix = numpy.zeros( - (interpList.shape[0], derivativeVector.shape[0], derivativeVector.shape[1]) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, derivativeVector.shape[0]): - derivativeMatrix[i, j] = derivativeVector[j].subs(x, interpList[i]) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - else: - raise NameError( - "derivativeMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - elif dim == 2: - x = Symbol("x") - y = Symbol("y") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [[0.5], [0.8660254037844385 * x], [0.8660254037844385 * y]] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.625 * x**2 * y**2 - 1.875 * y**2 - 1.875 * x**2 + 0.625], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - [6.5625 * x**4 - 5.625 * x**2 + 0.5625], - [6.5625 * y**4 - 5.625 * y**2 + 0.5625], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [[0.5], [0.8660254037844385 * x], [0.8660254037844385 * y], [1.5 * x * y]] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.625 * x**2 * y**2 - 1.875 * y**2 - 1.875 * x**2 + 0.625], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - [6.5625 * x**4 - 5.625 * x**2 + 0.5625], - [6.5625 * y**4 - 5.625 * y**2 + 0.5625], - [ - 11.36658342467074 * x**4 * y - - 9.74278579257492 * x**2 * y - + 0.9742785792574921 * y - ], - [ - 11.36658342467074 * x * y**4 - - 9.74278579257492 * x * y**2 - + 0.9742785792574921 * x - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [(x * y) / 4.0 - y / 4.0 - x / 4.0 + 1.0 / 4.0], - [x / 4.0 - y / 4.0 - (x * y) / 4.0 + 1.0 / 4.0], - [y / 4.0 - x / 4.0 - (x * y) / 4.0 + 1.0 / 4.0], - [x / 4.0 + y / 4.0 + (x * y) / 4.0 + 1.0 / 4.0], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [ - -(x**2 * y) / 4.0 - + x**2 / 4.0 - - (x * y**2) / 4.0 - + (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - [(x**2 * y) / 2.0 - y / 2.0 - x**2 / 2.0 + 1.0 / 2.0], - [ - -(x**2 * y) / 4.0 - + x**2 / 4.0 - + (x * y**2) / 4.0 - - (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - [(x * y**2) / 2.0 - x / 2.0 - y**2 / 2.0 + 1 / 2.0], - [x / 2.0 - (x * y**2) / 2.0 - y**2 / 2.0 + 1.0 / 2.0], - [ - (x**2 * y) / 4.0 - + x**2 / 4.0 - - (x * y**2) / 4.0 - - (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - [y / 2.0 - (x**2 * y) / 2.0 - x**2 / 2.0 + 1.0 / 2.0], - [ - (x**2 * y) / 4.0 - + x**2 / 4.0 - + (x * y**2) / 4.0 - + (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, derivativeVector.shape[0]): - for l in range(0, derivativeVector.shape[1]): - derivativeMatrix[j + i * interpList.shape[0], k, l] = ( - derivativeVector[k, l].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <3 for nodal Serendipity in 2D".format( - order - ) - ) - - else: - raise NameError( - "derivativeMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - elif dim == 3: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [ - 3.977475644174331 * x**2 * y**2 - - 1.325825214724777 * y**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * x**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * y**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * y**2 - + 0.4419417382415923 - ], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [ - 4.640388251536713 * x**4 - - 3.977475644174326 * x**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * y**4 - - 3.977475644174326 * y**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * z**4 - - 3.977475644174326 * z**2 - + 0.3977475644174325 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.837117307087383 * x * y * z], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [7.015607600201137 * x**3 * y * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y**3 * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y * z**3 - 4.209364560120682 * x * y * z], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [ - 3.977475644174331 * x**2 * y**2 - - 1.325825214724777 * y**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * x**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * y**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * y**2 - + 0.4419417382415923 - ], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [ - 4.640388251536713 * x**4 - - 3.977475644174326 * x**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * y**4 - - 3.977475644174326 * y**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * z**4 - - 3.977475644174326 * z**2 - + 0.3977475644174325 - ], - [ - 6.889189901577672 * x**2 * y**2 * z - - 2.296396633859224 * y**2 * z - - 2.296396633859224 * x**2 * z - + 0.7654655446197414 * z - ], - [ - 6.889189901577672 * x**2 * y * z**2 - - 2.296396633859224 * y * z**2 - - 2.296396633859224 * x**2 * y - + 0.7654655446197414 * y - ], - [ - 6.889189901577672 * x * y**2 * z**2 - - 2.296396633859224 * x * z**2 - - 2.296396633859224 * x * y**2 - + 0.7654655446197414 * x - ], - [7.015607600201137 * x**3 * y * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y**3 * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y * z**3 - 4.209364560120682 * x * y * z], - [ - 8.03738821850729 * x**4 * y - - 6.889189901577677 * x**2 * y - + 0.6889189901577677 * y - ], - [ - 8.03738821850729 * x * y**4 - - 6.889189901577677 * x * y**2 - + 0.6889189901577677 * x - ], - [ - 8.03738821850729 * x**4 * z - - 6.889189901577677 * x**2 * z - + 0.6889189901577677 * z - ], - [ - 8.03738821850729 * y**4 * z - - 6.889189901577677 * y**2 * z - + 0.6889189901577677 * z - ], - [ - 8.03738821850729 * x * z**4 - - 6.889189901577677 * x * z**2 - + 0.6889189901577677 * x - ], - [ - 8.03738821850729 * y * z**4 - - 6.889189901577677 * y * z**2 - + 0.6889189901577677 * y - ], - [ - 13.92116475461014 * x**4 * y * z - - 11.93242693252298 * x**2 * y * z - + 1.193242693252298 * y * z - ], - [ - 13.92116475461014 * x * y**4 * z - - 11.93242693252298 * x * y**2 * z - + 1.193242693252298 * x * z - ], - [ - 13.92116475461014 * x * y * z**4 - - 11.93242693252298 * x * y * z**2 - + 1.193242693252298 * x * y - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [ - (x * y) / 8.0 - - y / 8.0 - - z / 8.0 - - x / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - - z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - - z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - - z / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - z / 8.0 - - y / 8.0 - - x / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - + z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - + z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - + z / 8.0 - + (x * y) / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [ - (x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - (y * z) / 4.0 - - z / 4.0 - - y / 4.0 - + (x**2 * y) / 4.0 - + (x**2 * z) / 4.0 - - x**2 / 4.0 - - (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - (x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - (x * z) / 4.0 - - z / 4.0 - - x / 4.0 - + (x * y**2) / 4.0 - + (y**2 * z) / 4.0 - - y**2 / 4.0 - - (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - - z / 4.0 - - (x * z) / 4.0 - - (x * y**2) / 4.0 - + (y**2 * z) / 4.0 - - y**2 / 4.0 - + (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - y / 4.0 - - z / 4.0 - - (y * z) / 4.0 - - (x**2 * y) / 4.0 - + (x**2 * z) / 4.0 - - x**2 / 4.0 - + (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - (x * y) / 4.0 - - y / 4.0 - - x / 4.0 - + (x * z**2) / 4.0 - + (y * z**2) / 4.0 - - z**2 / 4.0 - - (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - - y / 4.0 - - (x * y) / 4.0 - - (x * z**2) / 4.0 - + (y * z**2) / 4.0 - - z**2 / 4.0 - + (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - y / 4.0 - - x / 4.0 - - (x * y) / 4.0 - + (x * z**2) / 4.0 - - (y * z**2) / 4.0 - - z**2 / 4.0 - + (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - + y / 4.0 - + (x * y) / 4.0 - - (x * z**2) / 4.0 - - (y * z**2) / 4.0 - - z**2 / 4.0 - - (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - [ - z / 4.0 - - y / 4.0 - - (y * z) / 4.0 - + (x**2 * y) / 4.0 - - (x**2 * z) / 4.0 - - x**2 / 4.0 - + (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - [ - z / 4.0 - - x / 4.0 - - (x * z) / 4.0 - + (x * y**2) / 4.0 - - (y**2 * z) / 4.0 - - y**2 / 4.0 - + (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - + z / 4.0 - + (x * z) / 4.0 - - (x * y**2) / 4.0 - - (y**2 * z) / 4.0 - - y**2 / 4.0 - - (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - (x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - [ - y / 4.0 - + z / 4.0 - + (y * z) / 4.0 - - (x**2 * y) / 4.0 - - (x**2 * z) / 4.0 - - x**2 / 4.0 - - (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - (x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, derivativeVector.shape[0]): - for m in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - m, - ] = ( - derivativeVector[l, m] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <3 for nodal Serendipity in 3D".format( - order - ) - ) - - else: - raise NameError( - "derivativeMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - elif dim == 4: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - w = Symbol("w") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [2.8125 * x**2 * y**2 - 0.9375 * y**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * x**2 * z**2 - 0.9375 * z**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * z**2 - 0.9375 * z**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * x**2 * w**2 - 0.9375 * w**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * w**2 - 0.9375 * w**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * z**2 * w**2 - 0.9375 * w**2 - 0.9375 * z**2 + 0.3125], - [2.864109809347398 * x**3 * y - 1.718465885608439 * x * y], - [2.864109809347398 * x * y**3 - 1.718465885608439 * x * y], - [2.864109809347398 * x**3 * z - 1.718465885608439 * x * z], - [2.864109809347398 * y**3 * z - 1.718465885608439 * y * z], - [2.864109809347398 * x * z**3 - 1.718465885608439 * x * z], - [2.864109809347398 * y * z**3 - 1.718465885608439 * y * z], - [2.864109809347398 * x**3 * w - 1.718465885608439 * x * w], - [2.864109809347398 * y**3 * w - 1.718465885608439 * y * w], - [2.864109809347398 * z**3 * w - 1.718465885608439 * z * w], - [2.864109809347398 * x * w**3 - 1.718465885608439 * x * w], - [2.864109809347398 * y * w**3 - 1.718465885608439 * y * w], - [2.864109809347398 * z * w**3 - 1.718465885608439 * z * w], - [3.28125 * x**4 - 2.8125 * x**2 + 0.28125], - [3.28125 * y**4 - 2.8125 * y**2 + 0.28125], - [3.28125 * z**4 - 2.8125 * z**2 + 0.28125], - [3.28125 * w**4 - 2.8125 * w**2 + 0.28125], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - elif modal and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [2.25 * x * y * z * w], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [2.864109809347398 * x**3 * y - 1.718465885608439 * x * y], - [2.864109809347398 * x * y**3 - 1.718465885608439 * x * y], - [2.864109809347398 * x**3 * z - 1.718465885608439 * x * z], - [2.864109809347398 * y**3 * z - 1.718465885608439 * y * z], - [2.864109809347398 * x * z**3 - 1.718465885608439 * x * z], - [2.864109809347398 * y * z**3 - 1.718465885608439 * y * z], - [2.864109809347398 * x**3 * w - 1.718465885608439 * x * w], - [2.864109809347398 * y**3 * w - 1.718465885608439 * y * w], - [2.864109809347398 * z**3 * w - 1.718465885608439 * z * w], - [2.864109809347398 * x * w**3 - 1.718465885608439 * x * w], - [2.864109809347398 * y * w**3 - 1.718465885608439 * y * w], - [2.864109809347398 * z * w**3 - 1.718465885608439 * z * w], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - [4.960783708246104 * x**3 * y * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y**3 * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y * z**3 - 2.976470224947662 * x * y * z], - [4.960783708246104 * x**3 * y * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * y**3 * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x**3 * z * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y**3 * z * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * z**3 * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z**3 * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * y * w**3 - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * z * w**3 - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z * w**3 - 2.976470224947662 * y * z * w], - [8.5923294280422 * x**3 * y * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y**3 * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z**3 * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z * w**3 - 5.15539765682532 * x * y * z * w], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [2.8125 * x**2 * y**2 - 0.9375 * y**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * x**2 * z**2 - 0.9375 * z**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * z**2 - 0.9375 * z**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * x**2 * w**2 - 0.9375 * w**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * w**2 - 0.9375 * w**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * z**2 * w**2 - 0.9375 * w**2 - 0.9375 * z**2 + 0.3125], - [2.864109809347398 * x**3 * y - 1.718465885608439 * x * y], - [2.864109809347398 * x * y**3 - 1.718465885608439 * x * y], - [2.864109809347398 * x**3 * z - 1.718465885608439 * x * z], - [2.864109809347398 * y**3 * z - 1.718465885608439 * y * z], - [2.864109809347398 * x * z**3 - 1.718465885608439 * x * z], - [2.864109809347398 * y * z**3 - 1.718465885608439 * y * z], - [2.864109809347398 * x**3 * w - 1.718465885608439 * x * w], - [2.864109809347398 * y**3 * w - 1.718465885608439 * y * w], - [2.864109809347398 * z**3 * w - 1.718465885608439 * z * w], - [2.864109809347398 * x * w**3 - 1.718465885608439 * x * w], - [2.864109809347398 * y * w**3 - 1.718465885608439 * y * w], - [2.864109809347398 * z * w**3 - 1.718465885608439 * z * w], - [3.28125 * x**4 - 2.8125 * x**2 + 0.28125], - [3.28125 * y**4 - 2.8125 * y**2 + 0.28125], - [3.28125 * z**4 - 2.8125 * z**2 + 0.28125], - [3.28125 * w**4 - 2.8125 * w**2 + 0.28125], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - [ - 4.87139289628746 * x**2 * y**2 * z - - 1.62379763209582 * y**2 * z - - 1.62379763209582 * x**2 * z - + 0.5412658773652733 * z - ], - [ - 4.87139289628746 * x**2 * y * z**2 - - 1.62379763209582 * y * z**2 - - 1.62379763209582 * x**2 * y - + 0.5412658773652733 * y - ], - [ - 4.87139289628746 * x * y**2 * z**2 - - 1.62379763209582 * x * z**2 - - 1.62379763209582 * x * y**2 - + 0.5412658773652733 * x - ], - [ - 4.87139289628746 * x**2 * y**2 * w - - 1.62379763209582 * y**2 * w - - 1.62379763209582 * x**2 * w - + 0.5412658773652733 * w - ], - [ - 4.87139289628746 * x**2 * z**2 * w - - 1.62379763209582 * z**2 * w - - 1.62379763209582 * x**2 * w - + 0.5412658773652733 * w - ], - [ - 4.87139289628746 * y**2 * z**2 * w - - 1.62379763209582 * z**2 * w - - 1.62379763209582 * y**2 * w - + 0.5412658773652733 * w - ], - [ - 4.87139289628746 * x**2 * y * w**2 - - 1.62379763209582 * y * w**2 - - 1.62379763209582 * x**2 * y - + 0.5412658773652733 * y - ], - [ - 4.87139289628746 * x * y**2 * w**2 - - 1.62379763209582 * x * w**2 - - 1.62379763209582 * x * y**2 - + 0.5412658773652733 * x - ], - [ - 4.87139289628746 * x**2 * z * w**2 - - 1.62379763209582 * z * w**2 - - 1.62379763209582 * x**2 * z - + 0.5412658773652733 * z - ], - [ - 4.87139289628746 * y**2 * z * w**2 - - 1.62379763209582 * z * w**2 - - 1.62379763209582 * y**2 * z - + 0.5412658773652733 * z - ], - [ - 4.87139289628746 * x * z**2 * w**2 - - 1.62379763209582 * x * w**2 - - 1.62379763209582 * x * z**2 - + 0.5412658773652733 * x - ], - [ - 4.87139289628746 * y * z**2 * w**2 - - 1.62379763209582 * y * w**2 - - 1.62379763209582 * y * z**2 - + 0.5412658773652733 * y - ], - [4.960783708246104 * x**3 * y * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y**3 * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y * z**3 - 2.976470224947662 * x * y * z], - [4.960783708246104 * x**3 * y * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * y**3 * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x**3 * z * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y**3 * z * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * z**3 * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z**3 * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * y * w**3 - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * z * w**3 - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z * w**3 - 2.976470224947662 * y * z * w], - [ - 5.68329171233537 * x**4 * y - - 4.87139289628746 * x**2 * y - + 0.487139289628746 * y - ], - [ - 5.68329171233537 * x * y**4 - - 4.87139289628746 * x * y**2 - + 0.487139289628746 * x - ], - [ - 5.68329171233537 * x**4 * z - - 4.87139289628746 * x**2 * z - + 0.487139289628746 * z - ], - [ - 5.68329171233537 * y**4 * z - - 4.87139289628746 * y**2 * z - + 0.487139289628746 * z - ], - [ - 5.68329171233537 * x * z**4 - - 4.87139289628746 * x * z**2 - + 0.487139289628746 * x - ], - [ - 5.68329171233537 * y * z**4 - - 4.87139289628746 * y * z**2 - + 0.487139289628746 * y - ], - [ - 5.68329171233537 * x**4 * w - - 4.87139289628746 * x**2 * w - + 0.487139289628746 * w - ], - [ - 5.68329171233537 * y**4 * w - - 4.87139289628746 * y**2 * w - + 0.487139289628746 * w - ], - [ - 5.68329171233537 * z**4 * w - - 4.87139289628746 * z**2 * w - + 0.487139289628746 * w - ], - [ - 5.68329171233537 * x * w**4 - - 4.87139289628746 * x * w**2 - + 0.487139289628746 * x - ], - [ - 5.68329171233537 * y * w**4 - - 4.87139289628746 * y * w**2 - + 0.487139289628746 * y - ], - [ - 5.68329171233537 * z * w**4 - - 4.87139289628746 * z * w**2 - + 0.487139289628746 * z - ], - [ - 8.4375 * x**2 * y**2 * z * w - - 2.8125 * y**2 * z * w - - 2.8125 * x**2 * z * w - + 0.9375 * z * w - ], - [ - 8.4375 * x**2 * y * z**2 * w - - 2.8125 * y * z**2 * w - - 2.8125 * x**2 * y * w - + 0.9375 * y * w - ], - [ - 8.4375 * x * y**2 * z**2 * w - - 2.8125 * x * z**2 * w - - 2.8125 * x * y**2 * w - + 0.9375 * x * w - ], - [ - 8.4375 * x**2 * y * z * w**2 - - 2.8125 * y * z * w**2 - - 2.8125 * x**2 * y * z - + 0.9375 * y * z - ], - [ - 8.4375 * x * y**2 * z * w**2 - - 2.8125 * x * z * w**2 - - 2.8125 * x * y**2 * z - + 0.9375 * x * z - ], - [ - 8.4375 * x * y * z**2 * w**2 - - 2.8125 * x * y * w**2 - - 2.8125 * x * y * z**2 - + 0.9375 * x * y - ], - [8.5923294280422 * x**3 * y * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y**3 * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z**3 * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z * w**3 - 5.15539765682532 * x * y * z * w], - [9.84375 * x**4 * y * z - 8.4375 * x**2 * y * z + 0.84375 * y * z], - [9.84375 * x * y**4 * z - 8.4375 * x * y**2 * z + 0.84375 * x * z], - [9.84375 * x * y * z**4 - 8.4375 * x * y * z**2 + 0.84375 * x * y], - [9.84375 * x**4 * y * w - 8.4375 * x**2 * y * w + 0.84375 * y * w], - [9.84375 * x * y**4 * w - 8.4375 * x * y**2 * w + 0.84375 * x * w], - [9.84375 * x**4 * z * w - 8.4375 * x**2 * z * w + 0.84375 * z * w], - [9.84375 * y**4 * z * w - 8.4375 * y**2 * z * w + 0.84375 * z * w], - [9.84375 * x * z**4 * w - 8.4375 * x * z**2 * w + 0.84375 * x * w], - [9.84375 * y * z**4 * w - 8.4375 * y * z**2 * w + 0.84375 * y * w], - [9.84375 * x * y * w**4 - 8.4375 * x * y * w**2 + 0.84375 * x * y], - [9.84375 * x * z * w**4 - 8.4375 * x * z * w**2 + 0.84375 * x * z], - [9.84375 * y * z * w**4 - 8.4375 * y * z * w**2 + 0.84375 * y * z], - [ - 17.04987513700614 * x**4 * y * z * w - - 14.61417868886241 * x**2 * y * z * w - + 1.46141786888624 * y * z * w - ], - [ - 17.04987513700614 * x * y**4 * z * w - - 14.61417868886241 * x * y**2 * z * w - + 1.46141786888624 * x * z * w - ], - [ - 17.04987513700614 * x * y * z**4 * w - - 14.61417868886241 * x * y * z**2 * w - + 1.46141786888624 * x * y * w - ], - [ - 17.04987513700614 * x * y * z * w**4 - - 14.61417868886241 * x * y * z * w**2 - + 1.46141786888624 * x * y * z - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [ - (w * x) / 16.0 - - x / 16.0 - - y / 16.0 - - z / 16.0 - - w / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - - y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - y / 16.0 - - x / 16.0 - - w / 16.0 - - z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - + y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - z / 16.0 - - x / 16.0 - - y / 16.0 - - w / 16.0 - + (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - - y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - y / 16.0 - - x / 16.0 - - w / 16.0 - + z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - + y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - - y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - - y / 16.0 - - z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - + y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - + y / 16.0 - - z / 16.0 - + (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - - y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - - y / 16.0 - + z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - + y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - + y / 16.0 - + z / 16.0 - + (w * x) / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - (w * y) / 8.0 - - y / 8.0 - - z / 8.0 - - w / 8.0 - + (w * z) / 8.0 - + (y * z) / 8.0 - + (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - (w * x) / 8.0 - - x / 8.0 - - z / 8.0 - - w / 8.0 - + (w * z) / 8.0 - + (x * z) / 8.0 - + (w * y**2) / 8.0 - + (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - - z / 8.0 - - (w * x) / 8.0 - + (w * z) / 8.0 - - (x * z) / 8.0 - + (w * y**2) / 8.0 - - (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - y / 8.0 - - w / 8.0 - - z / 8.0 - - (w * y) / 8.0 - + (w * z) / 8.0 - - (y * z) / 8.0 - + (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - (w * x) / 8.0 - - x / 8.0 - - y / 8.0 - - w / 8.0 - + (w * y) / 8.0 - + (x * y) / 8.0 - + (w * z**2) / 8.0 - + (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - - y / 8.0 - - (w * x) / 8.0 - + (w * y) / 8.0 - - (x * y) / 8.0 - + (w * z**2) / 8.0 - - (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - - w / 8.0 - + (w * x) / 8.0 - - (w * y) / 8.0 - - (x * y) / 8.0 - + (w * z**2) / 8.0 - + (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - + y / 8.0 - - (w * x) / 8.0 - - (w * y) / 8.0 - + (x * y) / 8.0 - + (w * z**2) / 8.0 - - (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - z / 8.0 - - y / 8.0 - - w / 8.0 - + (w * y) / 8.0 - - (w * z) / 8.0 - - (y * z) / 8.0 - + (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - z / 8.0 - - x / 8.0 - - w / 8.0 - + (w * x) / 8.0 - - (w * z) / 8.0 - - (x * z) / 8.0 - + (w * y**2) / 8.0 - + (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - + z / 8.0 - - (w * x) / 8.0 - - (w * z) / 8.0 - + (x * z) / 8.0 - + (w * y**2) / 8.0 - - (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - y / 8.0 - - w / 8.0 - + z / 8.0 - - (w * y) / 8.0 - - (w * z) / 8.0 - + (y * z) / 8.0 - + (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - (x * y) / 8.0 - - y / 8.0 - - z / 8.0 - - x / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - + (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - - z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - - (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - - z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - + (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - - z / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - - (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - z / 8.0 - - y / 8.0 - - x / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - + (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - + z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - - (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - + z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - + (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - + z / 8.0 - + (x * y) / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - - (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - y / 8.0 - - z / 8.0 - - (w * y) / 8.0 - - (w * z) / 8.0 - + (y * z) / 8.0 - - (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - x / 8.0 - - z / 8.0 - - (w * x) / 8.0 - - (w * z) / 8.0 - + (x * z) / 8.0 - - (w * y**2) / 8.0 - + (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - - z / 8.0 - + (w * x) / 8.0 - - (w * z) / 8.0 - - (x * z) / 8.0 - - (w * y**2) / 8.0 - - (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - + y / 8.0 - - z / 8.0 - + (w * y) / 8.0 - - (w * z) / 8.0 - - (y * z) / 8.0 - - (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - x / 8.0 - - y / 8.0 - - (w * x) / 8.0 - - (w * y) / 8.0 - + (x * y) / 8.0 - - (w * z**2) / 8.0 - + (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - - y / 8.0 - + (w * x) / 8.0 - - (w * y) / 8.0 - - (x * y) / 8.0 - - (w * z**2) / 8.0 - - (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - - x / 8.0 - + y / 8.0 - - (w * x) / 8.0 - + (w * y) / 8.0 - - (x * y) / 8.0 - - (w * z**2) / 8.0 - + (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - + y / 8.0 - + (w * x) / 8.0 - + (w * y) / 8.0 - + (x * y) / 8.0 - - (w * z**2) / 8.0 - - (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - y / 8.0 - + z / 8.0 - - (w * y) / 8.0 - + (w * z) / 8.0 - - (y * z) / 8.0 - - (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - x / 8.0 - + z / 8.0 - - (w * x) / 8.0 - + (w * z) / 8.0 - - (x * z) / 8.0 - - (w * y**2) / 8.0 - + (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - + z / 8.0 - + (w * x) / 8.0 - + (w * z) / 8.0 - + (x * z) / 8.0 - - (w * y**2) / 8.0 - - (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - + y / 8.0 - + z / 8.0 - + (w * y) / 8.0 - + (w * z) / 8.0 - + (y * z) / 8.0 - - (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, derivativeVector.shape[0]): - for n in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - n, - ] = ( - derivativeVector[m, n] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <3 for nodal Serendipity in 4D".format( - order - ) - ) - - else: - raise NameError( - "derivativeMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - - elif dim == 5: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - w = Symbol("w") - v = Symbol("v") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [ - 1.988737822087165 * x**2 * y**2 - - 0.6629126073623886 * y**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * w**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * w**2 - + 0.2209708691207962 - ], - [2.025231468252455 * x**3 * y - 1.215138880951473 * x * y], - [2.025231468252455 * x * y**3 - 1.215138880951473 * x * y], - [2.025231468252455 * x**3 * z - 1.215138880951473 * x * z], - [2.025231468252455 * y**3 * z - 1.215138880951473 * y * z], - [2.025231468252455 * x * z**3 - 1.215138880951473 * x * z], - [2.025231468252455 * y * z**3 - 1.215138880951473 * y * z], - [2.025231468252455 * x**3 * w - 1.215138880951473 * x * w], - [2.025231468252455 * y**3 * w - 1.215138880951473 * y * w], - [2.025231468252455 * z**3 * w - 1.215138880951473 * z * w], - [2.025231468252455 * x * w**3 - 1.215138880951473 * x * w], - [2.025231468252455 * y * w**3 - 1.215138880951473 * y * w], - [2.025231468252455 * z * w**3 - 1.215138880951473 * z * w], - [2.025231468252455 * x**3 * v - 1.215138880951473 * x * v], - [2.025231468252455 * y**3 * v - 1.215138880951473 * y * v], - [2.025231468252455 * z**3 * v - 1.215138880951473 * z * v], - [2.025231468252455 * w**3 * v - 1.215138880951473 * w * v], - [2.025231468252455 * x * v**3 - 1.215138880951473 * x * v], - [2.025231468252455 * y * v**3 - 1.215138880951473 * y * v], - [2.025231468252455 * z * v**3 - 1.215138880951473 * z * v], - [2.025231468252455 * w * v**3 - 1.215138880951473 * w * v], - [ - 2.320194125768356 * x**4 - - 1.988737822087163 * x**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * y**4 - - 1.988737822087163 * y**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * z**4 - - 1.988737822087163 * z**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * w**4 - - 1.988737822087163 * w**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * v**4 - - 1.988737822087163 * v**2 - + 0.1988737822087163 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [2.755675960631069 * x * y * z * w * v], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [2.755675960631069 * x * y * z * w * v], - [3.080939385966559 * x**2 * y * z * w - 1.026979795322186 * y * z * w], - [3.080939385966559 * x * y**2 * z * w - 1.026979795322186 * x * z * w], - [3.080939385966559 * x * y * z**2 * w - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * y * z * w**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x**2 * y * z * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y**2 * z * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * x * y * z**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x**2 * y * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * x**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * y**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * x * z**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * y * z**2 * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y * w**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x * z * w**2 * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * y * z * w**2 * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y * z * v**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x * y * w * v**2 - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * z * w * v**2 - 1.026979795322186 * x * z * w], - [3.080939385966559 * y * z * w * v**2 - 1.026979795322186 * y * z * w], - [ - 5.336343551534144 * x**2 * y * z * w * v - - 1.778781183844715 * y * z * w * v - ], - [ - 5.336343551534144 * x * y**2 * z * w * v - - 1.778781183844715 * x * z * w * v - ], - [ - 5.336343551534144 * x * y * z**2 * w * v - - 1.778781183844715 * x * y * w * v - ], - [ - 5.336343551534144 * x * y * z * w**2 * v - - 1.778781183844715 * x * y * z * v - ], - [ - 5.336343551534144 * x * y * z * w * v**2 - - 1.778781183844715 * x * y * z * w - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [2.025231468252455 * x**3 * y - 1.215138880951473 * x * y], - [2.025231468252455 * x * y**3 - 1.215138880951473 * x * y], - [2.025231468252455 * x**3 * z - 1.215138880951473 * x * z], - [2.025231468252455 * y**3 * z - 1.215138880951473 * y * z], - [2.025231468252455 * x * z**3 - 1.215138880951473 * x * z], - [2.025231468252455 * y * z**3 - 1.215138880951473 * y * z], - [2.025231468252455 * x**3 * w - 1.215138880951473 * x * w], - [2.025231468252455 * y**3 * w - 1.215138880951473 * y * w], - [2.025231468252455 * z**3 * w - 1.215138880951473 * z * w], - [2.025231468252455 * x * w**3 - 1.215138880951473 * x * w], - [2.025231468252455 * y * w**3 - 1.215138880951473 * y * w], - [2.025231468252455 * z * w**3 - 1.215138880951473 * z * w], - [2.025231468252455 * x**3 * v - 1.215138880951473 * x * v], - [2.025231468252455 * y**3 * v - 1.215138880951473 * y * v], - [2.025231468252455 * z**3 * v - 1.215138880951473 * z * v], - [2.025231468252455 * w**3 * v - 1.215138880951473 * w * v], - [2.025231468252455 * x * v**3 - 1.215138880951473 * x * v], - [2.025231468252455 * y * v**3 - 1.215138880951473 * y * v], - [2.025231468252455 * z * v**3 - 1.215138880951473 * z * v], - [2.025231468252455 * w * v**3 - 1.215138880951473 * w * v], - [2.755675960631069 * x * y * z * w * v], - [3.080939385966559 * x**2 * y * z * w - 1.026979795322186 * y * z * w], - [3.080939385966559 * x * y**2 * z * w - 1.026979795322186 * x * z * w], - [3.080939385966559 * x * y * z**2 * w - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * y * z * w**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x**2 * y * z * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y**2 * z * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * x * y * z**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x**2 * y * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * x**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * y**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * x * z**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * y * z**2 * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y * w**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x * z * w**2 * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * y * z * w**2 * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y * z * v**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x * y * w * v**2 - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * z * w * v**2 - 1.026979795322186 * x * z * w], - [3.080939385966559 * y * z * w * v**2 - 1.026979795322186 * y * z * w], - [3.507803800100568 * x**3 * y * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y**3 * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y * z**3 - 2.104682280060341 * x * y * z], - [3.507803800100568 * x**3 * y * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * y**3 * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x**3 * z * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y**3 * z * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * z**3 * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z**3 * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * y * w**3 - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * z * w**3 - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z * w**3 - 2.104682280060341 * y * z * w], - [3.507803800100568 * x**3 * y * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * y**3 * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x**3 * z * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y**3 * z * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * z**3 * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z**3 * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x**3 * w * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y**3 * w * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z**3 * w * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * w**3 * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w**3 * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w**3 * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * y * v**3 - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * z * v**3 - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z * v**3 - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * w * v**3 - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w * v**3 - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w * v**3 - 2.104682280060341 * z * w * v], - [ - 5.336343551534144 * x**2 * y * z * w * v - - 1.778781183844715 * y * z * w * v - ], - [ - 5.336343551534144 * x * y**2 * z * w * v - - 1.778781183844715 * x * z * w * v - ], - [ - 5.336343551534144 * x * y * z**2 * w * v - - 1.778781183844715 * x * y * w * v - ], - [ - 5.336343551534144 * x * y * z * w**2 * v - - 1.778781183844715 * x * y * z * v - ], - [ - 5.336343551534144 * x * y * z * w * v**2 - - 1.778781183844715 * x * y * z * w - ], - [ - 6.075694404757367 * x**3 * y * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y**3 * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z**3 * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z * w**3 - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x**3 * y * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y**3 * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * z**3 * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x**3 * y * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * y**3 * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x**3 * z * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y**3 * z * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * z**3 * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z**3 * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * w**3 * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w**3 * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w**3 * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * z * v**3 - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * w * v**3 - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w * v**3 - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w * v**3 - - 3.64541664285442 * y * z * w * v - ], - [ - 10.52341140030171 * x**3 * y * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y**3 * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z**3 * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w**3 * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w * v**3 - - 6.314046840181025 * x * y * z * w * v - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [ - 1.988737822087165 * x**2 * y**2 - - 0.6629126073623886 * y**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * w**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * w**2 - + 0.2209708691207962 - ], - [2.025231468252455 * x**3 * y - 1.215138880951473 * x * y], - [2.025231468252455 * x * y**3 - 1.215138880951473 * x * y], - [2.025231468252455 * x**3 * z - 1.215138880951473 * x * z], - [2.025231468252455 * y**3 * z - 1.215138880951473 * y * z], - [2.025231468252455 * x * z**3 - 1.215138880951473 * x * z], - [2.025231468252455 * y * z**3 - 1.215138880951473 * y * z], - [2.025231468252455 * x**3 * w - 1.215138880951473 * x * w], - [2.025231468252455 * y**3 * w - 1.215138880951473 * y * w], - [2.025231468252455 * z**3 * w - 1.215138880951473 * z * w], - [2.025231468252455 * x * w**3 - 1.215138880951473 * x * w], - [2.025231468252455 * y * w**3 - 1.215138880951473 * y * w], - [2.025231468252455 * z * w**3 - 1.215138880951473 * z * w], - [2.025231468252455 * x**3 * v - 1.215138880951473 * x * v], - [2.025231468252455 * y**3 * v - 1.215138880951473 * y * v], - [2.025231468252455 * z**3 * v - 1.215138880951473 * z * v], - [2.025231468252455 * w**3 * v - 1.215138880951473 * w * v], - [2.025231468252455 * x * v**3 - 1.215138880951473 * x * v], - [2.025231468252455 * y * v**3 - 1.215138880951473 * y * v], - [2.025231468252455 * z * v**3 - 1.215138880951473 * z * v], - [2.025231468252455 * w * v**3 - 1.215138880951473 * w * v], - [ - 2.320194125768356 * x**4 - - 1.988737822087163 * x**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * y**4 - - 1.988737822087163 * y**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * z**4 - - 1.988737822087163 * z**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * w**4 - - 1.988737822087163 * w**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * v**4 - - 1.988737822087163 * v**2 - + 0.1988737822087163 - ], - [2.755675960631069 * x * y * z * w * v], - [3.080939385966559 * x**2 * y * z * w - 1.026979795322186 * y * z * w], - [3.080939385966559 * x * y**2 * z * w - 1.026979795322186 * x * z * w], - [3.080939385966559 * x * y * z**2 * w - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * y * z * w**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x**2 * y * z * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y**2 * z * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * x * y * z**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x**2 * y * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * x**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * y**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * x * z**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * y * z**2 * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y * w**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x * z * w**2 * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * y * z * w**2 * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y * z * v**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x * y * w * v**2 - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * z * w * v**2 - 1.026979795322186 * x * z * w], - [3.080939385966559 * y * z * w * v**2 - 1.026979795322186 * y * z * w], - [ - 3.444594950788842 * x**2 * y**2 * z - - 1.148198316929614 * y**2 * z - - 1.148198316929614 * x**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * x**2 * y * z**2 - - 1.148198316929614 * y * z**2 - - 1.148198316929614 * x**2 * y - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x * y**2 * z**2 - - 1.148198316929614 * x * z**2 - - 1.148198316929614 * x * y**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * x**2 * y**2 * w - - 1.148198316929614 * y**2 * w - - 1.148198316929614 * x**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * x**2 * z**2 * w - - 1.148198316929614 * z**2 * w - - 1.148198316929614 * x**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * y**2 * z**2 * w - - 1.148198316929614 * z**2 * w - - 1.148198316929614 * y**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * x**2 * y * w**2 - - 1.148198316929614 * y * w**2 - - 1.148198316929614 * x**2 * y - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x * y**2 * w**2 - - 1.148198316929614 * x * w**2 - - 1.148198316929614 * x * y**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * x**2 * z * w**2 - - 1.148198316929614 * z * w**2 - - 1.148198316929614 * x**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * y**2 * z * w**2 - - 1.148198316929614 * z * w**2 - - 1.148198316929614 * y**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * x * z**2 * w**2 - - 1.148198316929614 * x * w**2 - - 1.148198316929614 * x * z**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * y * z**2 * w**2 - - 1.148198316929614 * y * w**2 - - 1.148198316929614 * y * z**2 - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x**2 * y**2 * v - - 1.148198316929614 * y**2 * v - - 1.148198316929614 * x**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * x**2 * z**2 * v - - 1.148198316929614 * z**2 * v - - 1.148198316929614 * x**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * y**2 * z**2 * v - - 1.148198316929614 * z**2 * v - - 1.148198316929614 * y**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * x**2 * w**2 * v - - 1.148198316929614 * w**2 * v - - 1.148198316929614 * x**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * y**2 * w**2 * v - - 1.148198316929614 * w**2 * v - - 1.148198316929614 * y**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * z**2 * w**2 * v - - 1.148198316929614 * w**2 * v - - 1.148198316929614 * z**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * x**2 * y * v**2 - - 1.148198316929614 * y * v**2 - - 1.148198316929614 * x**2 * y - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x * y**2 * v**2 - - 1.148198316929614 * x * v**2 - - 1.148198316929614 * x * y**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * x**2 * z * v**2 - - 1.148198316929614 * z * v**2 - - 1.148198316929614 * x**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * y**2 * z * v**2 - - 1.148198316929614 * z * v**2 - - 1.148198316929614 * y**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * x * z**2 * v**2 - - 1.148198316929614 * x * v**2 - - 1.148198316929614 * x * z**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * y * z**2 * v**2 - - 1.148198316929614 * y * v**2 - - 1.148198316929614 * y * z**2 - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x**2 * w * v**2 - - 1.148198316929614 * w * v**2 - - 1.148198316929614 * x**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * y**2 * w * v**2 - - 1.148198316929614 * w * v**2 - - 1.148198316929614 * y**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * z**2 * w * v**2 - - 1.148198316929614 * w * v**2 - - 1.148198316929614 * z**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * x * w**2 * v**2 - - 1.148198316929614 * x * v**2 - - 1.148198316929614 * x * w**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * y * w**2 * v**2 - - 1.148198316929614 * y * v**2 - - 1.148198316929614 * y * w**2 - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * z * w**2 * v**2 - - 1.148198316929614 * z * v**2 - - 1.148198316929614 * z * w**2 - + 0.3827327723098713 * z - ], - [3.507803800100568 * x**3 * y * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y**3 * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y * z**3 - 2.104682280060341 * x * y * z], - [3.507803800100568 * x**3 * y * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * y**3 * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x**3 * z * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y**3 * z * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * z**3 * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z**3 * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * y * w**3 - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * z * w**3 - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z * w**3 - 2.104682280060341 * y * z * w], - [3.507803800100568 * x**3 * y * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * y**3 * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x**3 * z * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y**3 * z * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * z**3 * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z**3 * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x**3 * w * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y**3 * w * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z**3 * w * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * w**3 * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w**3 * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w**3 * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * y * v**3 - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * z * v**3 - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z * v**3 - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * w * v**3 - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w * v**3 - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w * v**3 - 2.104682280060341 * z * w * v], - [ - 4.018694109253645 * x**4 * y - - 3.444594950788839 * x**2 * y - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * x * y**4 - - 3.444594950788839 * x * y**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * x**4 * z - - 3.444594950788839 * x**2 * z - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * y**4 * z - - 3.444594950788839 * y**2 * z - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * x * z**4 - - 3.444594950788839 * x * z**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * y * z**4 - - 3.444594950788839 * y * z**2 - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * x**4 * w - - 3.444594950788839 * x**2 * w - + 0.3444594950788838 * w - ], - [ - 4.018694109253645 * y**4 * w - - 3.444594950788839 * y**2 * w - + 0.3444594950788838 * w - ], - [ - 4.018694109253645 * z**4 * w - - 3.444594950788839 * z**2 * w - + 0.3444594950788838 * w - ], - [ - 4.018694109253645 * x * w**4 - - 3.444594950788839 * x * w**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * y * w**4 - - 3.444594950788839 * y * w**2 - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * z * w**4 - - 3.444594950788839 * z * w**2 - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * x**4 * v - - 3.444594950788839 * x**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * y**4 * v - - 3.444594950788839 * y**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * z**4 * v - - 3.444594950788839 * z**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * w**4 * v - - 3.444594950788839 * w**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * x * v**4 - - 3.444594950788839 * x * v**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * y * v**4 - - 3.444594950788839 * y * v**2 - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * z * v**4 - - 3.444594950788839 * z * v**2 - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * w * v**4 - - 3.444594950788839 * w * v**2 - + 0.3444594950788838 * w - ], - [ - 5.336343551534144 * x**2 * y * z * w * v - - 1.778781183844715 * y * z * w * v - ], - [ - 5.336343551534144 * x * y**2 * z * w * v - - 1.778781183844715 * x * z * w * v - ], - [ - 5.336343551534144 * x * y * z**2 * w * v - - 1.778781183844715 * x * y * w * v - ], - [ - 5.336343551534144 * x * y * z * w**2 * v - - 1.778781183844715 * x * y * z * v - ], - [ - 5.336343551534144 * x * y * z * w * v**2 - - 1.778781183844715 * x * y * z * w - ], - [ - 5.966213466261497 * x**2 * y**2 * z * w - - 1.988737822087165 * y**2 * z * w - - 1.988737822087165 * x**2 * z * w - + 0.6629126073623886 * z * w - ], - [ - 5.966213466261497 * x**2 * y * z**2 * w - - 1.988737822087165 * y * z**2 * w - - 1.988737822087165 * x**2 * y * w - + 0.6629126073623886 * y * w - ], - [ - 5.966213466261497 * x * y**2 * z**2 * w - - 1.988737822087165 * x * z**2 * w - - 1.988737822087165 * x * y**2 * w - + 0.6629126073623886 * x * w - ], - [ - 5.966213466261497 * x**2 * y * z * w**2 - - 1.988737822087165 * y * z * w**2 - - 1.988737822087165 * x**2 * y * z - + 0.6629126073623886 * y * z - ], - [ - 5.966213466261497 * x * y**2 * z * w**2 - - 1.988737822087165 * x * z * w**2 - - 1.988737822087165 * x * y**2 * z - + 0.6629126073623886 * x * z - ], - [ - 5.966213466261497 * x * y * z**2 * w**2 - - 1.988737822087165 * x * y * w**2 - - 1.988737822087165 * x * y * z**2 - + 0.6629126073623886 * x * y - ], - [ - 5.966213466261497 * x**2 * y**2 * z * v - - 1.988737822087165 * y**2 * z * v - - 1.988737822087165 * x**2 * z * v - + 0.6629126073623886 * z * v - ], - [ - 5.966213466261497 * x**2 * y * z**2 * v - - 1.988737822087165 * y * z**2 * v - - 1.988737822087165 * x**2 * y * v - + 0.6629126073623886 * y * v - ], - [ - 5.966213466261497 * x * y**2 * z**2 * v - - 1.988737822087165 * x * z**2 * v - - 1.988737822087165 * x * y**2 * v - + 0.6629126073623886 * x * v - ], - [ - 5.966213466261497 * x**2 * y**2 * w * v - - 1.988737822087165 * y**2 * w * v - - 1.988737822087165 * x**2 * w * v - + 0.6629126073623886 * w * v - ], - [ - 5.966213466261497 * x**2 * z**2 * w * v - - 1.988737822087165 * z**2 * w * v - - 1.988737822087165 * x**2 * w * v - + 0.6629126073623886 * w * v - ], - [ - 5.966213466261497 * y**2 * z**2 * w * v - - 1.988737822087165 * z**2 * w * v - - 1.988737822087165 * y**2 * w * v - + 0.6629126073623886 * w * v - ], - [ - 5.966213466261497 * x**2 * y * w**2 * v - - 1.988737822087165 * y * w**2 * v - - 1.988737822087165 * x**2 * y * v - + 0.6629126073623886 * y * v - ], - [ - 5.966213466261497 * x * y**2 * w**2 * v - - 1.988737822087165 * x * w**2 * v - - 1.988737822087165 * x * y**2 * v - + 0.6629126073623886 * x * v - ], - [ - 5.966213466261497 * x**2 * z * w**2 * v - - 1.988737822087165 * z * w**2 * v - - 1.988737822087165 * x**2 * z * v - + 0.6629126073623886 * z * v - ], - [ - 5.966213466261497 * y**2 * z * w**2 * v - - 1.988737822087165 * z * w**2 * v - - 1.988737822087165 * y**2 * z * v - + 0.6629126073623886 * z * v - ], - [ - 5.966213466261497 * x * z**2 * w**2 * v - - 1.988737822087165 * x * w**2 * v - - 1.988737822087165 * x * z**2 * v - + 0.6629126073623886 * x * v - ], - [ - 5.966213466261497 * y * z**2 * w**2 * v - - 1.988737822087165 * y * w**2 * v - - 1.988737822087165 * y * z**2 * v - + 0.6629126073623886 * y * v - ], - [ - 5.966213466261497 * x**2 * y * z * v**2 - - 1.988737822087165 * y * z * v**2 - - 1.988737822087165 * x**2 * y * z - + 0.6629126073623886 * y * z - ], - [ - 5.966213466261497 * x * y**2 * z * v**2 - - 1.988737822087165 * x * z * v**2 - - 1.988737822087165 * x * y**2 * z - + 0.6629126073623886 * x * z - ], - [ - 5.966213466261497 * x * y * z**2 * v**2 - - 1.988737822087165 * x * y * v**2 - - 1.988737822087165 * x * y * z**2 - + 0.6629126073623886 * x * y - ], - [ - 5.966213466261497 * x**2 * y * w * v**2 - - 1.988737822087165 * y * w * v**2 - - 1.988737822087165 * x**2 * y * w - + 0.6629126073623886 * y * w - ], - [ - 5.966213466261497 * x * y**2 * w * v**2 - - 1.988737822087165 * x * w * v**2 - - 1.988737822087165 * x * y**2 * w - + 0.6629126073623886 * x * w - ], - [ - 5.966213466261497 * x**2 * z * w * v**2 - - 1.988737822087165 * z * w * v**2 - - 1.988737822087165 * x**2 * z * w - + 0.6629126073623886 * z * w - ], - [ - 5.966213466261497 * y**2 * z * w * v**2 - - 1.988737822087165 * z * w * v**2 - - 1.988737822087165 * y**2 * z * w - + 0.6629126073623886 * z * w - ], - [ - 5.966213466261497 * x * z**2 * w * v**2 - - 1.988737822087165 * x * w * v**2 - - 1.988737822087165 * x * z**2 * w - + 0.6629126073623886 * x * w - ], - [ - 5.966213466261497 * y * z**2 * w * v**2 - - 1.988737822087165 * y * w * v**2 - - 1.988737822087165 * y * z**2 * w - + 0.6629126073623886 * y * w - ], - [ - 5.966213466261497 * x * y * w**2 * v**2 - - 1.988737822087165 * x * y * v**2 - - 1.988737822087165 * x * y * w**2 - + 0.6629126073623886 * x * y - ], - [ - 5.966213466261497 * x * z * w**2 * v**2 - - 1.988737822087165 * x * z * v**2 - - 1.988737822087165 * x * z * w**2 - + 0.6629126073623886 * x * z - ], - [ - 5.966213466261497 * y * z * w**2 * v**2 - - 1.988737822087165 * y * z * v**2 - - 1.988737822087165 * y * z * w**2 - + 0.6629126073623886 * y * z - ], - [ - 6.075694404757367 * x**3 * y * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y**3 * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z**3 * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z * w**3 - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x**3 * y * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y**3 * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * z**3 * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x**3 * y * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * y**3 * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x**3 * z * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y**3 * z * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * z**3 * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z**3 * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * w**3 * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w**3 * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w**3 * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * z * v**3 - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * w * v**3 - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w * v**3 - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w * v**3 - - 3.64541664285442 * y * z * w * v - ], - [ - 6.960582377305069 * x**4 * y * z - - 5.966213466261488 * x**2 * y * z - + 0.5966213466261489 * y * z - ], - [ - 6.960582377305069 * x * y**4 * z - - 5.966213466261488 * x * y**2 * z - + 0.5966213466261489 * x * z - ], - [ - 6.960582377305069 * x * y * z**4 - - 5.966213466261488 * x * y * z**2 - + 0.5966213466261489 * x * y - ], - [ - 6.960582377305069 * x**4 * y * w - - 5.966213466261488 * x**2 * y * w - + 0.5966213466261489 * y * w - ], - [ - 6.960582377305069 * x * y**4 * w - - 5.966213466261488 * x * y**2 * w - + 0.5966213466261489 * x * w - ], - [ - 6.960582377305069 * x**4 * z * w - - 5.966213466261488 * x**2 * z * w - + 0.5966213466261489 * z * w - ], - [ - 6.960582377305069 * y**4 * z * w - - 5.966213466261488 * y**2 * z * w - + 0.5966213466261489 * z * w - ], - [ - 6.960582377305069 * x * z**4 * w - - 5.966213466261488 * x * z**2 * w - + 0.5966213466261489 * x * w - ], - [ - 6.960582377305069 * y * z**4 * w - - 5.966213466261488 * y * z**2 * w - + 0.5966213466261489 * y * w - ], - [ - 6.960582377305069 * x * y * w**4 - - 5.966213466261488 * x * y * w**2 - + 0.5966213466261489 * x * y - ], - [ - 6.960582377305069 * x * z * w**4 - - 5.966213466261488 * x * z * w**2 - + 0.5966213466261489 * x * z - ], - [ - 6.960582377305069 * y * z * w**4 - - 5.966213466261488 * y * z * w**2 - + 0.5966213466261489 * y * z - ], - [ - 6.960582377305069 * x**4 * y * v - - 5.966213466261488 * x**2 * y * v - + 0.5966213466261489 * y * v - ], - [ - 6.960582377305069 * x * y**4 * v - - 5.966213466261488 * x * y**2 * v - + 0.5966213466261489 * x * v - ], - [ - 6.960582377305069 * x**4 * z * v - - 5.966213466261488 * x**2 * z * v - + 0.5966213466261489 * z * v - ], - [ - 6.960582377305069 * y**4 * z * v - - 5.966213466261488 * y**2 * z * v - + 0.5966213466261489 * z * v - ], - [ - 6.960582377305069 * x * z**4 * v - - 5.966213466261488 * x * z**2 * v - + 0.5966213466261489 * x * v - ], - [ - 6.960582377305069 * y * z**4 * v - - 5.966213466261488 * y * z**2 * v - + 0.5966213466261489 * y * v - ], - [ - 6.960582377305069 * x**4 * w * v - - 5.966213466261488 * x**2 * w * v - + 0.5966213466261489 * w * v - ], - [ - 6.960582377305069 * y**4 * w * v - - 5.966213466261488 * y**2 * w * v - + 0.5966213466261489 * w * v - ], - [ - 6.960582377305069 * z**4 * w * v - - 5.966213466261488 * z**2 * w * v - + 0.5966213466261489 * w * v - ], - [ - 6.960582377305069 * x * w**4 * v - - 5.966213466261488 * x * w**2 * v - + 0.5966213466261489 * x * v - ], - [ - 6.960582377305069 * y * w**4 * v - - 5.966213466261488 * y * w**2 * v - + 0.5966213466261489 * y * v - ], - [ - 6.960582377305069 * z * w**4 * v - - 5.966213466261488 * z * w**2 * v - + 0.5966213466261489 * z * v - ], - [ - 6.960582377305069 * x * y * v**4 - - 5.966213466261488 * x * y * v**2 - + 0.5966213466261489 * x * y - ], - [ - 6.960582377305069 * x * z * v**4 - - 5.966213466261488 * x * z * v**2 - + 0.5966213466261489 * x * z - ], - [ - 6.960582377305069 * y * z * v**4 - - 5.966213466261488 * y * z * v**2 - + 0.5966213466261489 * y * z - ], - [ - 6.960582377305069 * x * w * v**4 - - 5.966213466261488 * x * w * v**2 - + 0.5966213466261489 * x * w - ], - [ - 6.960582377305069 * y * w * v**4 - - 5.966213466261488 * y * w * v**2 - + 0.5966213466261489 * y * w - ], - [ - 6.960582377305069 * z * w * v**4 - - 5.966213466261488 * z * w * v**2 - + 0.5966213466261489 * z * w - ], - [ - 10.33378485236653 * x**2 * y**2 * z * w * v - - 3.444594950788842 * y**2 * z * w * v - - 3.444594950788842 * x**2 * z * w * v - + 1.148198316929614 * z * w * v - ], - [ - 10.33378485236653 * x**2 * y * z**2 * w * v - - 3.444594950788842 * y * z**2 * w * v - - 3.444594950788842 * x**2 * y * w * v - + 1.148198316929614 * y * w * v - ], - [ - 10.33378485236653 * x * y**2 * z**2 * w * v - - 3.444594950788842 * x * z**2 * w * v - - 3.444594950788842 * x * y**2 * w * v - + 1.148198316929614 * x * w * v - ], - [ - 10.33378485236653 * x**2 * y * z * w**2 * v - - 3.444594950788842 * y * z * w**2 * v - - 3.444594950788842 * x**2 * y * z * v - + 1.148198316929614 * y * z * v - ], - [ - 10.33378485236653 * x * y**2 * z * w**2 * v - - 3.444594950788842 * x * z * w**2 * v - - 3.444594950788842 * x * y**2 * z * v - + 1.148198316929614 * x * z * v - ], - [ - 10.33378485236653 * x * y * z**2 * w**2 * v - - 3.444594950788842 * x * y * w**2 * v - - 3.444594950788842 * x * y * z**2 * v - + 1.148198316929614 * x * y * v - ], - [ - 10.33378485236653 * x**2 * y * z * w * v**2 - - 3.444594950788842 * y * z * w * v**2 - - 3.444594950788842 * x**2 * y * z * w - + 1.148198316929614 * y * z * w - ], - [ - 10.33378485236653 * x * y**2 * z * w * v**2 - - 3.444594950788842 * x * z * w * v**2 - - 3.444594950788842 * x * y**2 * z * w - + 1.148198316929614 * x * z * w - ], - [ - 10.33378485236653 * x * y * z**2 * w * v**2 - - 3.444594950788842 * x * y * w * v**2 - - 3.444594950788842 * x * y * z**2 * w - + 1.148198316929614 * x * y * w - ], - [ - 10.33378485236653 * x * y * z * w**2 * v**2 - - 3.444594950788842 * x * y * z * v**2 - - 3.444594950788842 * x * y * z * w**2 - + 1.148198316929614 * x * y * z - ], - [ - 10.52341140030171 * x**3 * y * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y**3 * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z**3 * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w**3 * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w * v**3 - - 6.314046840181025 * x * y * z * w * v - ], - [ - 12.05608232776096 * x**4 * y * z * w - - 10.33378485236654 * x**2 * y * z * w - + 1.033378485236654 * y * z * w - ], - [ - 12.05608232776096 * x * y**4 * z * w - - 10.33378485236654 * x * y**2 * z * w - + 1.033378485236654 * x * z * w - ], - [ - 12.05608232776096 * x * y * z**4 * w - - 10.33378485236654 * x * y * z**2 * w - + 1.033378485236654 * x * y * w - ], - [ - 12.05608232776096 * x * y * z * w**4 - - 10.33378485236654 * x * y * z * w**2 - + 1.033378485236654 * x * y * z - ], - [ - 12.05608232776096 * x**4 * y * z * v - - 10.33378485236654 * x**2 * y * z * v - + 1.033378485236654 * y * z * v - ], - [ - 12.05608232776096 * x * y**4 * z * v - - 10.33378485236654 * x * y**2 * z * v - + 1.033378485236654 * x * z * v - ], - [ - 12.05608232776096 * x * y * z**4 * v - - 10.33378485236654 * x * y * z**2 * v - + 1.033378485236654 * x * y * v - ], - [ - 12.05608232776096 * x**4 * y * w * v - - 10.33378485236654 * x**2 * y * w * v - + 1.033378485236654 * y * w * v - ], - [ - 12.05608232776096 * x * y**4 * w * v - - 10.33378485236654 * x * y**2 * w * v - + 1.033378485236654 * x * w * v - ], - [ - 12.05608232776096 * x**4 * z * w * v - - 10.33378485236654 * x**2 * z * w * v - + 1.033378485236654 * z * w * v - ], - [ - 12.05608232776096 * y**4 * z * w * v - - 10.33378485236654 * y**2 * z * w * v - + 1.033378485236654 * z * w * v - ], - [ - 12.05608232776096 * x * z**4 * w * v - - 10.33378485236654 * x * z**2 * w * v - + 1.033378485236654 * x * w * v - ], - [ - 12.05608232776096 * y * z**4 * w * v - - 10.33378485236654 * y * z**2 * w * v - + 1.033378485236654 * y * w * v - ], - [ - 12.05608232776096 * x * y * w**4 * v - - 10.33378485236654 * x * y * w**2 * v - + 1.033378485236654 * x * y * v - ], - [ - 12.05608232776096 * x * z * w**4 * v - - 10.33378485236654 * x * z * w**2 * v - + 1.033378485236654 * x * z * v - ], - [ - 12.05608232776096 * y * z * w**4 * v - - 10.33378485236654 * y * z * w**2 * v - + 1.033378485236654 * y * z * v - ], - [ - 12.05608232776096 * x * y * z * v**4 - - 10.33378485236654 * x * y * z * v**2 - + 1.033378485236654 * x * y * z - ], - [ - 12.05608232776096 * x * y * w * v**4 - - 10.33378485236654 * x * y * w * v**2 - + 1.033378485236654 * x * y * w - ], - [ - 12.05608232776096 * x * z * w * v**4 - - 10.33378485236654 * x * z * w * v**2 - + 1.033378485236654 * x * z * w - ], - [ - 12.05608232776096 * y * z * w * v**4 - - 10.33378485236654 * y * z * w * v**2 - + 1.033378485236654 * y * z * w - ], - [ - 20.88174713191521 * x**4 * y * z * w * v - - 17.89864039878447 * x**2 * y * z * w * v - + 1.789864039878446 * y * z * w * v - ], - [ - 20.88174713191521 * x * y**4 * z * w * v - - 17.89864039878447 * x * y**2 * z * w * v - + 1.789864039878446 * x * z * w * v - ], - [ - 20.88174713191521 * x * y * z**4 * w * v - - 17.89864039878447 * x * y * z**2 * w * v - + 1.789864039878446 * x * y * w * v - ], - [ - 20.88174713191521 * x * y * z * w**4 * v - - 17.89864039878447 * x * y * z * w**2 * v - + 1.789864039878446 * x * y * z * v - ], - [ - 20.88174713191521 * x * y * z * w * v**4 - - 17.89864039878447 * x * y * z * w * v**2 - + 1.789864039878446 * x * y * z * w - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [ - (v * w) / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - - v / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - - y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - y / 32.0 - - w / 32.0 - - x / 32.0 - - v / 32.0 - - z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - + y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - z / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - - v / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - - y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - y / 32.0 - - w / 32.0 - - x / 32.0 - - v / 32.0 - + z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - + y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - - y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - + y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - + y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - - y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - - y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - + y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - + y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - ] - ) - - derivativeVector = zeros(functionVector.shape[0], dim) - - for i in range(0, functionVector.shape[0]): - derivativeVector[i, 0] = diff(functionVector[i], x) - derivativeVector[i, 1] = diff(functionVector[i], y) - derivativeVector[i, 2] = diff(functionVector[i], z) - derivativeVector[i, 3] = diff(functionVector[i], w) - derivativeVector[i, 4] = diff(functionVector[i], v) - - derivativeMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - derivativeVector.shape[0], - derivativeVector.shape[1], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, derivativeVector.shape[0]): - for o in range(0, derivativeVector.shape[1]): - derivativeMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - o, - ] = ( - derivativeVector[n, o] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - else: - raise NameError( - "derivativeMatrix: Order {} is not supported!\nPolynomial order must be 1 for nodal Serendipity in 5D".format( - order - ) - ) - - else: - raise NameError( - "derivativeMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - - else: - raise NameError("derivativeMatrix: Dimension {} is not supported.".format(dim)) - - return derivativeMatrix - - -if __name__ == "__main__": - import tables - # set command line options - parser = OptionParser() - parser.add_option( - "-d", "--dimension", action="store", dest="dim", help="specified dimension" - ) - parser.add_option( - "-o", "--order", action="store", dest="order", help="specified polynomial order" - ) - parser.add_option( - "-b", "--basis", action="store", dest="basis", help="specified basis set" - ) - parser.add_option( - "-i", - "--interp", - action="store", - dest="interp", - help="specified number of interpolation points", - ) - parser.add_option( - "-m", - "--modal", - action="store", - dest="modal", - help="set to True for modal basis set", - ) - - (options, args) = parser.parse_args() - - dim = int(options.dim) - order = int(options.order) - basis_type = options.basis - modal = options.modal - interp = int(options.interp) - - derivativeMatrix = createDerivativeMatrix(dim, order, basis_type, interp, modal) - fh = tables.open_file("derivativeMatrix.h5", mode="w") - fh.create_array("/", "derivative_matrix", derivativeMatrix) - fh.close() diff --git a/src/postgkyl/data/computeInterpolationMatrices.py b/src/postgkyl/data/computeInterpolationMatrices.py deleted file mode 100644 index 2912c510..00000000 --- a/src/postgkyl/data/computeInterpolationMatrices.py +++ /dev/null @@ -1,9064 +0,0 @@ -import numpy -from sympy import * - -from optparse import OptionParser - - -def createInterpMatrix(dim, order, basis_type, interp, modal=True, c2p=False): - if c2p: - interp += 1 - # end - interpList = numpy.zeros(interp) - for i in range(interp): - if c2p: - interpList[i] = -1.0 + float(i) * 2.0 / (interp - 1) - else: - interpList[i] = -1.0 * (interp - 1) / interp + float(i) * 2.0 / interp - # end - # end - - # The following is for gkhybrid only. - interpListND = list() - for d in range(dim): - interp_true = interp - if basis_type == "gkhybrid": - # 1x1v, 1x2v, 2x2v, 3x2v cases, with p=2 in the first velocity dim. - if ( - ((dim == 2 or dim == 3) and d == 1) - or (dim == 4 and d == 2) - or (dim == 5 and d == 3) - ): - interp_true = interp + 1 - # end - elif basis_type == "gkhybrid_vel": - # 1v, 2v, with p=2 in the first velocity dim. - if (d == 0): - interp_true = interp + 1 - # end - elif basis_type == "hybrid": - # 1x1v, 2x2v, 2x2v, 3x2v cases, with p=2 in the first velocity dim. - if d == dim - 1: - interp_true = interp + 1 - # end - # end - - interpListND.append(numpy.zeros(interp_true)) - for i in range(interp_true): - if c2p: - interpListND[d][i] = -1.0 + float(i) * 2.0 / (interp_true - 1) - else: - interpListND[d][i] = ( - -1.0 * (interp_true - 1) / interp_true + float(i) * 2.0 / interp_true - ) - # end - # end - # end - - if dim == 1: - x = Symbol("x") - if modal and basis_type == "gkhybrid_vel": - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - ] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif modal: - if order == 0: - functionVector = Matrix([[0.7071067811865468]]) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 1: - functionVector = Matrix([[0.7071067811865468], [1.224744871391589 * x]]) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 2: - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - ] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 3: - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - [4.677071733467427 * x**3 - 2.806243040080457 * x], - ] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 4: - functionVector = Matrix( - [ - [0.7071067811865468], - [1.224744871391589 * x], - [2.371708245126285 * x**2 - 0.7905694150420951], - [4.677071733467427 * x**3 - 2.806243040080457 * x], - [ - 9.280776503073431 * x**4 - - 7.954951288348656 * x**2 - + 0.7954951288348655 - ], - ] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - # end - else: - if order == 1: - functionVector = Matrix([[0.5 - 0.5 * x], [0.5 + 0.5 * x]]) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 2: - functionVector = Matrix( - [[0.5 * x**2 - 0.5 * x], [1.0 - x**2], [0.5 * x**2 + 0.5 * x]] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 3: - functionVector = Matrix( - [ - [-(9.0 * x**3) / 16.0 + (9.0 * x**2) / 16.0 + x / 16.0 - 1 / 16.0], - [ - (27.0 * x**3) / 16.0 - - (9.0 * x**2) / 16.0 - - (27.0 * x) / 16.0 - + 9.0 / 16.0 - ], - [ - (27.0 * x) / 16.0 - - (9.0 * x**2) / 16.0 - - (27.0 * x**3) / 16.0 - + 9.0 / 16.0 - ], - [(9.0 * x**3) / 16.0 + (9.0 * x**2) / 16.0 - x / 16.0 - 1 / 16.0], - ] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - elif order == 4: - functionVector = Matrix( - [ - [(2.0 * x**4) / 3.0 - (2.0 * x**3) / 3.0 - x**2 / 6.0 + x / 6.0], - [ - -(8.0 * x**4) / 3.0 - + (4.0 * x**3) / 3.0 - + (8.0 * x**2) / 3.0 - - (4.0 * x) / 3.0 - ], - [4.0 * x**4 - 5.0 * x**2 + 1.0], - [ - -(8.0 * x**4) / 3.0 - - (4.0 * x**3) / 3.0 - + (8.0 * x**2) / 3.0 - + (4.0 * x) / 3.0 - ], - [(2.0 * x**4) / 3.0 + (2.0 * x**3) / 3.0 - x**2 / 6.0 - x / 6.0], - ] - ) - interpMatrix = numpy.zeros((interpList.shape[0], functionVector.shape[0])) - for i in range(0, interpList.shape[0]): - for j in range(0, functionVector.shape[0]): - interpMatrix[i, j] = functionVector[j].subs(x, interpList[i]) - # end - # end - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - # end - # end - elif dim == 2: - x = Symbol("x") - y = Symbol("y") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [[0.5], [0.8660254037844385 * x], [0.8660254037844385 * y]] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.625 * x**2 * y**2 - 1.875 * y**2 - 1.875 * x**2 + 0.625], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - [6.5625 * x**4 - 5.625 * x**2 + 0.5625], - [6.5625 * y**4 - 5.625 * y**2 + 0.5625], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "serendipity": - if order == 0: - functionVector = Matrix([[0.5]]) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - elif order == 1: - functionVector = Matrix( - [[0.5], [0.8660254037844385 * x], [0.8660254037844385 * y], [1.5 * x * y]] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.625 * x**2 * y**2 - 1.875 * y**2 - 1.875 * x**2 + 0.625], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - [6.5625 * x**4 - 5.625 * x**2 + 0.5625], - [6.5625 * y**4 - 5.625 * y**2 + 0.5625], - [ - 11.36658342467074 * x**4 * y - - 9.74278579257492 * x**2 * y - + 0.9742785792574921 * y - ], - [ - 11.36658342467074 * x * y**4 - - 9.74278579257492 * x * y**2 - + 0.9742785792574921 * x - ], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "tensor": - if order == 1: - functionVector = Matrix( - [[0.5], [0.8660254037844385 * x], [0.8660254037844385 * y], [1.5 * x * y]] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [5.625 * x**2 * y**2 - 1.875 * y**2 - 1.875 * x**2 + 0.625], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844385 * x], - [0.8660254037844385 * y], - [1.5 * x * y], - [1.677050983124845 * x**2 - 0.5590169943749485], - [1.677050983124845 * y**2 - 0.5590169943749485], - [2.904737509655563 * x**2 * y - 0.9682458365518544 * y], - [2.904737509655563 * x * y**2 - 0.9682458365518544 * x], - [3.307189138830737 * x**3 - 1.984313483298442 * x], - [3.307189138830737 * y**3 - 1.984313483298442 * y], - [5.625 * x**2 * y**2 - 1.875 * y**2 - 1.875 * x**2 + 0.625], - [5.728219618694792 * x**3 * y - 3.436931771216875 * x * y], - [5.728219618694792 * x * y**3 - 3.436931771216875 * x * y], - [ - 11.09264959331178 * x**3 * y**2 - - 6.655589755987068 * x * y**2 - - 3.69754986443726 * x**3 - + 2.218529918662355 * x - ], - [ - 11.09264959331178 * x**2 * y**3 - - 3.69754986443726 * y**3 - - 6.655589755987068 * x**2 * y - + 2.218529918662355 * y - ], - [ - 21.875 * x**3 * y**3 - - 13.125 * x * y**3 - - 13.125 * x**3 * y - + 7.875 * x * y - ], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <4".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [(x * y) / 4.0 - y / 4.0 - x / 4.0 + 1.0 / 4.0], - [x / 4.0 - y / 4.0 - (x * y) / 4.0 + 1.0 / 4.0], - [y / 4.0 - x / 4.0 - (x * y) / 4.0 + 1.0 / 4.0], - [x / 4.0 + y / 4.0 + (x * y) / 4.0 + 1.0 / 4.0], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [ - -(x**2 * y) / 4.0 - + x**2 / 4.0 - - (x * y**2) / 4.0 - + (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - [(x**2 * y) / 2.0 - y / 2.0 - x**2 / 2.0 + 1.0 / 2.0], - [ - -(x**2 * y) / 4.0 - + x**2 / 4.0 - + (x * y**2) / 4.0 - - (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - [(x * y**2) / 2.0 - x / 2.0 - y**2 / 2.0 + 1 / 2.0], - [x / 2.0 - (x * y**2) / 2.0 - y**2 / 2.0 + 1.0 / 2.0], - [ - (x**2 * y) / 4.0 - + x**2 / 4.0 - - (x * y**2) / 4.0 - - (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - [y / 2.0 - (x**2 * y) / 2.0 - x**2 / 2.0 + 1.0 / 2.0], - [ - (x**2 * y) / 4.0 - + x**2 / 4.0 - + (x * y**2) / 4.0 - + (x * y) / 4.0 - + y**2 / 4.0 - - 1 / 4.0 - ], - ] - ) - interpMatrix = numpy.zeros( - (interpList.shape[0] * interpList.shape[0], functionVector.shape[0]) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpList.shape[0], k] = ( - functionVector[k].subs(x, interpList[j]).subs(y, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <3 for nodal Serendipity in 2D".format( - order - ) - ) - - elif modal and basis_type == "gkhybrid": - if order == 1: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844386 * x], - [0.8660254037844386 * y], - [1.5 * x * y], - [1.677050983124842 * (y**2 - 0.3333333333333333)], - [2.904737509655563 * (x * y**2 - 0.3333333333333333 * x)], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] * interpListND[1].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[1].shape[0]): - for j in range(0, interpListND[0].shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpListND[0].shape[0], k] = ( - functionVector[k] - .subs(x, interpListND[0][j]) - .subs(y, interpListND[1][i]) - ) - - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal and basis_type == "gkhybrid_vel": - if order == 1: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844386 * x], - [0.8660254037844386 * y], - [1.5 * x * y], - [1.677050983124842 * (x**2 - 0.3333333333333333)], - [2.904737509655563 * (x**2 * y- 0.3333333333333333 * y)], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] * interpListND[1].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[1].shape[0]): - for j in range(0, interpListND[0].shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpListND[0].shape[0], k] = ( - functionVector[k] - .subs(x, interpListND[0][j]) - .subs(y, interpListND[1][i]) - ) - - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal and basis_type == "hybrid": - if order == 1: - functionVector = Matrix( - [ - [0.5], - [0.8660254037844386 * x], - [0.8660254037844386 * y], - [1.5 * x * y], - [1.677050983124842 * (y**2 - 0.3333333333333333)], - [2.904737509655563 * (x * y**2 - 0.3333333333333333 * x)], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] * interpListND[1].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[1].shape[0]): - for j in range(0, interpListND[0].shape[0]): - for k in range(0, functionVector.shape[0]): - interpMatrix[j + i * interpListND[0].shape[0], k] = ( - functionVector[k] - .subs(x, interpListND[0][j]) - .subs(y, interpListND[1][i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - else: - raise NameError( - "interpMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - elif dim == 3: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [ - 3.977475644174331 * x**2 * y**2 - - 1.325825214724777 * y**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * x**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * y**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * y**2 - + 0.4419417382415923 - ], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [ - 4.640388251536713 * x**4 - - 3.977475644174326 * x**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * y**4 - - 3.977475644174326 * y**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * z**4 - - 3.977475644174326 * z**2 - + 0.3977475644174325 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "serendipity": - if order == 0: - functionVector = Matrix([[0.3535533905932734]]) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - elif order == 1: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.837117307087383 * x * y * z], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [7.015607600201137 * x**3 * y * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y**3 * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y * z**3 - 4.209364560120682 * x * y * z], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [ - 3.977475644174331 * x**2 * y**2 - - 1.325825214724777 * y**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * x**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * x**2 - + 0.4419417382415923 - ], - [ - 3.977475644174331 * y**2 * z**2 - - 1.325825214724777 * z**2 - - 1.325825214724777 * y**2 - + 0.4419417382415923 - ], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [ - 4.640388251536713 * x**4 - - 3.977475644174326 * x**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * y**4 - - 3.977475644174326 * y**2 - + 0.3977475644174325 - ], - [ - 4.640388251536713 * z**4 - - 3.977475644174326 * z**2 - + 0.3977475644174325 - ], - [ - 6.889189901577672 * x**2 * y**2 * z - - 2.296396633859224 * y**2 * z - - 2.296396633859224 * x**2 * z - + 0.7654655446197414 * z - ], - [ - 6.889189901577672 * x**2 * y * z**2 - - 2.296396633859224 * y * z**2 - - 2.296396633859224 * x**2 * y - + 0.7654655446197414 * y - ], - [ - 6.889189901577672 * x * y**2 * z**2 - - 2.296396633859224 * x * z**2 - - 2.296396633859224 * x * y**2 - + 0.7654655446197414 * x - ], - [7.015607600201137 * x**3 * y * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y**3 * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y * z**3 - 4.209364560120682 * x * y * z], - [ - 8.03738821850729 * x**4 * y - - 6.889189901577677 * x**2 * y - + 0.6889189901577677 * y - ], - [ - 8.03738821850729 * x * y**4 - - 6.889189901577677 * x * y**2 - + 0.6889189901577677 * x - ], - [ - 8.03738821850729 * x**4 * z - - 6.889189901577677 * x**2 * z - + 0.6889189901577677 * z - ], - [ - 8.03738821850729 * y**4 * z - - 6.889189901577677 * y**2 * z - + 0.6889189901577677 * z - ], - [ - 8.03738821850729 * x * z**4 - - 6.889189901577677 * x * z**2 - + 0.6889189901577677 * x - ], - [ - 8.03738821850729 * y * z**4 - - 6.889189901577677 * y * z**2 - + 0.6889189901577677 * y - ], - [ - 13.92116475461014 * x**4 * y * z - - 11.93242693252298 * x**2 * y * z - + 1.193242693252298 * y * z - ], - [ - 13.92116475461014 * x * y**4 * z - - 11.93242693252298 * x * y**2 * z - + 1.193242693252298 * x * z - ], - [ - 13.92116475461014 * x * y * z**4 - - 11.93242693252298 * x * y * z**2 - + 1.193242693252298 * x * y - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "tensor": - if order == 1: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.837117307087383 * x * y * z], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [ - 3.977475644174328 * x**2 * y**2 - - 1.325825214724776 * y**2 - - 1.325825214724776 * x**2 - + 0.441941738241592 - ], - [ - 3.977475644174328 * x**2 * z**2 - - 1.325825214724776 * z**2 - - 1.325825214724776 * x**2 - + 0.441941738241592 - ], - [ - 3.977475644174328 * y**2 * z**2 - - 1.325825214724776 * z**2 - - 1.325825214724776 * y**2 - + 0.441941738241592 - ], - [ - 6.889189901577683 * x**2 * y**2 * z - - 2.296396633859227 * y**2 * z - - 2.296396633859227 * x**2 * z - + 0.7654655446197425 * z - ], - [ - 6.889189901577683 * x**2 * y * z**2 - - 2.296396633859227 * y * z**2 - - 2.296396633859227 * x**2 * y - + 0.7654655446197425 * y - ], - [ - 6.889189901577683 * x * y**2 * z**2 - - 2.296396633859227 * x * z**2 - - 2.296396633859227 * x * y**2 - + 0.7654655446197425 * x - ], - [ - 13.34085887883535 * x**2 * y**2 * z**2 - - 4.446952959611782 * y**2 * z**2 - - 4.446952959611782 * x**2 * z**2 - + 1.482317653203927 * z**2 - - 4.446952959611782 * x**2 * y**2 - + 1.482317653203927 * y**2 - + 1.482317653203927 * x**2 - - 0.4941058844013091 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.3535533905932734], - [0.6123724356957931 * x], - [0.6123724356957931 * y], - [0.6123724356957931 * z], - [1.060660171779822 * x * y], - [1.060660171779822 * x * z], - [1.060660171779822 * y * z], - [1.185854122563141 * x**2 - 0.3952847075210471], - [1.185854122563141 * y**2 - 0.3952847075210471], - [1.185854122563141 * z**2 - 0.3952847075210471], - [1.837117307087383 * x * y * z], - [2.053959590644373 * x**2 * y - 0.6846531968814578 * y], - [2.053959590644373 * x * y**2 - 0.6846531968814578 * x], - [2.053959590644373 * x**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * y**2 * z - 0.6846531968814578 * z], - [2.053959590644373 * x * z**2 - 0.6846531968814578 * x], - [2.053959590644373 * y * z**2 - 0.6846531968814578 * y], - [2.338535866733713 * x**3 - 1.403121520040228 * x], - [2.338535866733713 * y**3 - 1.403121520040228 * y], - [2.338535866733713 * z**3 - 1.403121520040228 * z], - [3.557562367689424 * x**2 * y * z - 1.185854122563141 * y * z], - [3.557562367689424 * x * y**2 * z - 1.185854122563141 * x * z], - [3.557562367689424 * x * y * z**2 - 1.185854122563141 * x * y], - [ - 3.977475644174328 * x**2 * y**2 - - 1.325825214724776 * y**2 - - 1.325825214724776 * x**2 - + 0.441941738241592 - ], - [ - 3.977475644174328 * x**2 * z**2 - - 1.325825214724776 * z**2 - - 1.325825214724776 * x**2 - + 0.441941738241592 - ], - [ - 3.977475644174328 * y**2 * z**2 - - 1.325825214724776 * z**2 - - 1.325825214724776 * y**2 - + 0.441941738241592 - ], - [4.050462936504911 * x**3 * y - 2.430277761902947 * x * y], - [4.050462936504911 * x * y**3 - 2.430277761902947 * x * y], - [4.050462936504911 * x**3 * z - 2.430277761902947 * x * z], - [4.050462936504911 * y**3 * z - 2.430277761902947 * y * z], - [4.050462936504911 * x * z**3 - 2.430277761902947 * x * z], - [4.050462936504911 * y * z**3 - 2.430277761902947 * y * z], - [ - 6.889189901577683 * x**2 * y**2 * z - - 2.296396633859227 * y**2 * z - - 2.296396633859227 * x**2 * z - + 0.7654655446197425 * z - ], - [ - 6.889189901577683 * x**2 * y * z**2 - - 2.296396633859227 * y * z**2 - - 2.296396633859227 * x**2 * y - + 0.7654655446197425 * y - ], - [ - 6.889189901577683 * x * y**2 * z**2 - - 2.296396633859227 * x * z**2 - - 2.296396633859227 * x * y**2 - + 0.7654655446197425 * x - ], - [7.015607600201137 * x**3 * y * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y**3 * z - 4.209364560120682 * x * y * z], - [7.015607600201137 * x * y * z**3 - 4.209364560120682 * x * y * z], - [ - 7.843687748756954 * x**3 * y**2 - - 4.706212649254172 * x * y**2 - - 2.614562582918984 * x**3 - + 1.56873754975139 * x - ], - [ - 7.843687748756954 * x**2 * y**3 - - 2.614562582918984 * y**3 - - 4.706212649254172 * x**2 * y - + 1.56873754975139 * y - ], - [ - 7.843687748756954 * x**3 * z**2 - - 4.706212649254172 * x * z**2 - - 2.614562582918984 * x**3 - + 1.56873754975139 * x - ], - [ - 7.843687748756954 * y**3 * z**2 - - 4.706212649254172 * y * z**2 - - 2.614562582918984 * y**3 - + 1.56873754975139 * y - ], - [ - 7.843687748756954 * x**2 * z**3 - - 2.614562582918984 * z**3 - - 4.706212649254172 * x**2 * z - + 1.56873754975139 * z - ], - [ - 7.843687748756954 * y**2 * z**3 - - 2.614562582918984 * z**3 - - 4.706212649254172 * y**2 * z - + 1.56873754975139 * z - ], - [ - 13.34085887883535 * x**2 * y**2 * z**2 - - 4.446952959611782 * y**2 * z**2 - - 4.446952959611782 * x**2 * z**2 - + 1.482317653203927 * z**2 - - 4.446952959611782 * x**2 * y**2 - + 1.482317653203927 * y**2 - + 1.482317653203927 * x**2 - - 0.4941058844013091 - ], - [ - 13.58566569955259 * x**3 * y**2 * z - - 8.151399419731556 * x * y**2 * z - - 4.528555233184197 * x**3 * z - + 2.717133139910518 * x * z - ], - [ - 13.58566569955259 * x**2 * y**3 * z - - 4.528555233184197 * y**3 * z - - 8.151399419731556 * x**2 * y * z - + 2.717133139910518 * y * z - ], - [ - 13.58566569955259 * x**3 * y * z**2 - - 8.151399419731556 * x * y * z**2 - - 4.528555233184197 * x**3 * y - + 2.717133139910518 * x * y - ], - [ - 13.58566569955259 * x * y**3 * z**2 - - 8.151399419731556 * x * y * z**2 - - 4.528555233184197 * x * y**3 - + 2.717133139910518 * x * y - ], - [ - 13.58566569955259 * x**2 * y * z**3 - - 4.528555233184197 * y * z**3 - - 8.151399419731556 * x**2 * y * z - + 2.717133139910518 * y * z - ], - [ - 13.58566569955259 * x * y**2 * z**3 - - 4.528555233184197 * x * z**3 - - 8.151399419731556 * x * y**2 * z - + 2.717133139910518 * x * z - ], - [ - 15.46796083845572 * x**3 * y**3 - - 9.280776503073431 * x * y**3 - - 9.280776503073431 * x**3 * y - + 5.568465901844059 * x * y - ], - [ - 15.46796083845572 * x**3 * z**3 - - 9.280776503073431 * x * z**3 - - 9.280776503073431 * x**3 * z - + 5.568465901844059 * x * z - ], - [ - 15.46796083845572 * y**3 * z**3 - - 9.280776503073431 * y * z**3 - - 9.280776503073431 * y**3 * z - + 5.568465901844059 * y * z - ], - [ - 26.30852850075426 * x**3 * y**2 * z**2 - - 15.78511710045256 * x * y**2 * z**2 - - 8.76950950025142 * x**3 * z**2 - + 5.261705700150851 * x * z**2 - - 8.76950950025142 * x**3 * y**2 - + 5.261705700150851 * x * y**2 - + 2.92316983341714 * x**3 - - 1.753901900050284 * x - ], - [ - 26.30852850075426 * x**2 * y**3 * z**2 - - 8.76950950025142 * y**3 * z**2 - - 15.78511710045256 * x**2 * y * z**2 - + 5.261705700150851 * y * z**2 - - 8.76950950025142 * x**2 * y**3 - + 2.92316983341714 * y**3 - + 5.261705700150851 * x**2 * y - - 1.753901900050284 * y - ], - [ - 26.30852850075426 * x**2 * y**2 * z**3 - - 8.76950950025142 * y**2 * z**3 - - 8.76950950025142 * x**2 * z**3 - + 2.92316983341714 * z**3 - - 15.78511710045256 * x**2 * y**2 * z - + 5.261705700150851 * y**2 * z - + 5.261705700150851 * x**2 * z - - 1.753901900050284 * z - ], - [ - 26.791294061691 * x**3 * y**3 * z - - 16.0747764370146 * x * y**3 * z - - 16.0747764370146 * x**3 * y * z - + 9.644865862208759 * x * y * z - ], - [ - 26.791294061691 * x**3 * y * z**3 - - 16.0747764370146 * x * y * z**3 - - 16.0747764370146 * x**3 * y * z - + 9.644865862208759 * x * y * z - ], - [ - 26.791294061691 * x * y**3 * z**3 - - 16.0747764370146 * x * y * z**3 - - 16.0747764370146 * x * y**3 * z - + 9.644865862208759 * x * y * z - ], - [ - 51.88111786213746 * x**3 * y**3 * z**2 - - 31.12867071728247 * x * y**3 * z**2 - - 31.12867071728247 * x**3 * y * z**2 - + 18.67720243036948 * x * y * z**2 - - 17.29370595404582 * x**3 * y**3 - + 10.37622357242749 * x * y**3 - + 10.37622357242749 * x**3 * y - - 6.225734143456492 * x * y - ], - [ - 51.88111786213746 * x**3 * y**2 * z**3 - - 31.12867071728247 * x * y**2 * z**3 - - 17.29370595404582 * x**3 * z**3 - + 10.37622357242749 * x * z**3 - - 31.12867071728247 * x**3 * y**2 * z - + 18.67720243036948 * x * y**2 * z - + 10.37622357242749 * x**3 * z - - 6.225734143456492 * x * z - ], - [ - 51.88111786213746 * x**2 * y**3 * z**3 - - 17.29370595404582 * y**3 * z**3 - - 31.12867071728247 * x**2 * y * z**3 - + 10.37622357242749 * y * z**3 - - 31.12867071728247 * x**2 * y**3 * z - + 10.37622357242749 * y**3 * z - + 18.67720243036948 * x**2 * y * z - - 6.225734143456492 * y * z - ], - [ - 102.3109441695999 * x**3 * y**3 * z**3 - - 61.38656650175994 * x * y**3 * z**3 - - 61.38656650175994 * x**3 * y * z**3 - + 36.83193990105597 * x * y * z**3 - - 61.38656650175994 * x**3 * y**3 * z - + 36.83193990105597 * x * y**3 * z - + 36.83193990105597 * x**3 * y * z - - 22.09916394063358 * x * y * z - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <4".format( - order - ) - ) - - elif modal and basis_type == "gkhybrid": - if order == 1: - functionVector = Matrix( - [ - [0.3535533905932737], - [0.6123724356957944 * x], - [0.6123724356957944 * y], - [0.6123724356957944 * z], - [1.060660171779821 * x * y], - [1.060660171779821 * x * z], - [1.060660171779821 * y * z], - [1.837117307087383 * x * y * z], - [1.185854122563142 * (y**2 - 0.3333333333333333)], - [2.053959590644372 * (x * y**2 - 0.3333333333333333 * x)], - [2.053959590644372 * (y**2 * z - 0.3333333333333333 * z)], - [3.557562367689425 * (x * y**2 * z - 0.3333333333333333 * x * z)], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[2].shape[0]): - for j in range(0, interpListND[1].shape[0]): - for k in range(0, interpListND[0].shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpListND[0].shape[0] - + i * interpListND[1].shape[0] * interpListND[0].shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpListND[0][k]) - .subs(y, interpListND[1][j]) - .subs(z, interpListND[2][i]) - ) - - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal and basis_type == "hybrid": - if order == 1: - functionVector = Matrix( - [ - [0.3535533905932737], - [0.6123724356957945 * x], - [0.6123724356957945 * y], - [0.6123724356957945 * z], - [1.060660171779821 * x * y], - [1.060660171779821 * x * z], - [1.060660171779821 * y * z], - [1.837117307087384 * x * y * z], - [1.185854122563142 * (z**2 - 0.3333333333333333)], - [2.053959590644373 * (x * z**2 - 0.3333333333333333 * x)], - [2.053959590644373 * (y * z**2 - 0.3333333333333333 * y)], - [3.557562367689427 * (x * y * z**2 - 0.3333333333333332 * x * y)], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[2].shape[0]): - for j in range(0, interpListND[1].shape[0]): - for k in range(0, interpListND[0].shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpListND[0].shape[0] - + i * interpListND[1].shape[0] * interpListND[0].shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpListND[0][k]) - .subs(y, interpListND[1][j]) - .subs(z, interpListND[2][i]) - ) - - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [ - (x * y) / 8.0 - - y / 8.0 - - z / 8.0 - - x / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - - z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - - z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - - z / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - z / 8.0 - - y / 8.0 - - x / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - + z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - + z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - - (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - + z / 8.0 - + (x * y) / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - + (x * y * z) / 8.0 - + 1.0 / 8.0 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [ - (x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - (y * z) / 4.0 - - z / 4.0 - - y / 4.0 - + (x**2 * y) / 4.0 - + (x**2 * z) / 4.0 - - x**2 / 4.0 - - (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - (x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - (x * z) / 4.0 - - z / 4.0 - - x / 4.0 - + (x * y**2) / 4.0 - + (y**2 * z) / 4.0 - - y**2 / 4.0 - - (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - - z / 4.0 - - (x * z) / 4.0 - - (x * y**2) / 4.0 - + (y**2 * z) / 4.0 - - y**2 / 4.0 - + (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - y / 4.0 - - z / 4.0 - - (y * z) / 4.0 - - (x**2 * y) / 4.0 - + (x**2 * z) / 4.0 - - x**2 / 4.0 - + (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - - (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - + z / 8.0 - - 1.0 / 4.0 - ], - [ - (x * y) / 4.0 - - y / 4.0 - - x / 4.0 - + (x * z**2) / 4.0 - + (y * z**2) / 4.0 - - z**2 / 4.0 - - (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - - y / 4.0 - - (x * y) / 4.0 - - (x * z**2) / 4.0 - + (y * z**2) / 4.0 - - z**2 / 4.0 - + (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - y / 4.0 - - x / 4.0 - - (x * y) / 4.0 - + (x * z**2) / 4.0 - - (y * z**2) / 4.0 - - z**2 / 4.0 - + (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - + y / 4.0 - + (x * y) / 4.0 - - (x * z**2) / 4.0 - - (y * z**2) / 4.0 - - z**2 / 4.0 - - (x * y * z**2) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - [ - z / 4.0 - - y / 4.0 - - (y * z) / 4.0 - + (x**2 * y) / 4.0 - - (x**2 * z) / 4.0 - - x**2 / 4.0 - + (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - -(x**2 * y * z) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - - (y * z**2) / 8.0 - + y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - [ - z / 4.0 - - x / 4.0 - - (x * z) / 4.0 - + (x * y**2) / 4.0 - - (y**2 * z) / 4.0 - - y**2 / 4.0 - + (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - x / 4.0 - + z / 4.0 - + (x * z) / 4.0 - - (x * y**2) / 4.0 - - (y**2 * z) / 4.0 - - y**2 / 4.0 - - (x * y**2 * z) / 4.0 - + 1.0 / 4.0 - ], - [ - (x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - - (x * y**2 * z) / 8.0 - - (x * y**2) / 8.0 - - (x * y * z**2) / 8.0 - - (x * y * z) / 8.0 - - (x * z**2) / 8.0 - + x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - [ - y / 4.0 - + z / 4.0 - + (y * z) / 4.0 - - (x**2 * y) / 4.0 - - (x**2 * z) / 4.0 - - x**2 / 4.0 - - (x**2 * y * z) / 4.0 - + 1.0 / 4.0 - ], - [ - (x**2 * y * z) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - + x**2 / 8.0 - + (x * y**2 * z) / 8.0 - + (x * y**2) / 8.0 - + (x * y * z**2) / 8.0 - + (x * y * z) / 8.0 - + (x * z**2) / 8.0 - - x / 8.0 - + (y**2 * z) / 8.0 - + y**2 / 8.0 - + (y * z**2) / 8.0 - - y / 8.0 - + z**2 / 8.0 - - z / 8.0 - - 1.0 / 4.0 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] * interpList.shape[0] * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, functionVector.shape[0]): - interpMatrix[ - k - + j * interpList.shape[0] - + i * interpList.shape[0] * interpList.shape[0], - l, - ] = ( - functionVector[l] - .subs(x, interpList[k]) - .subs(y, interpList[j]) - .subs(z, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <3 for nodal Serendipity in 3D".format( - order - ) - ) - - else: - raise NameError( - "interpMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - elif dim == 4: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - w = Symbol("w") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [2.8125 * x**2 * y**2 - 0.9375 * y**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * x**2 * z**2 - 0.9375 * z**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * z**2 - 0.9375 * z**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * x**2 * w**2 - 0.9375 * w**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * w**2 - 0.9375 * w**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * z**2 * w**2 - 0.9375 * w**2 - 0.9375 * z**2 + 0.3125], - [2.864109809347398 * x**3 * y - 1.718465885608439 * x * y], - [2.864109809347398 * x * y**3 - 1.718465885608439 * x * y], - [2.864109809347398 * x**3 * z - 1.718465885608439 * x * z], - [2.864109809347398 * y**3 * z - 1.718465885608439 * y * z], - [2.864109809347398 * x * z**3 - 1.718465885608439 * x * z], - [2.864109809347398 * y * z**3 - 1.718465885608439 * y * z], - [2.864109809347398 * x**3 * w - 1.718465885608439 * x * w], - [2.864109809347398 * y**3 * w - 1.718465885608439 * y * w], - [2.864109809347398 * z**3 * w - 1.718465885608439 * z * w], - [2.864109809347398 * x * w**3 - 1.718465885608439 * x * w], - [2.864109809347398 * y * w**3 - 1.718465885608439 * y * w], - [2.864109809347398 * z * w**3 - 1.718465885608439 * z * w], - [3.28125 * x**4 - 2.8125 * x**2 + 0.28125], - [3.28125 * y**4 - 2.8125 * y**2 + 0.28125], - [3.28125 * z**4 - 2.8125 * z**2 + 0.28125], - [3.28125 * w**4 - 2.8125 * w**2 + 0.28125], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - elif modal and basis_type == "serendipity": - if order == 0: - functionVector = Matrix([[0.25]]) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [2.25 * x * y * z * w], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [2.864109809347398 * x**3 * y - 1.718465885608439 * x * y], - [2.864109809347398 * x * y**3 - 1.718465885608439 * x * y], - [2.864109809347398 * x**3 * z - 1.718465885608439 * x * z], - [2.864109809347398 * y**3 * z - 1.718465885608439 * y * z], - [2.864109809347398 * x * z**3 - 1.718465885608439 * x * z], - [2.864109809347398 * y * z**3 - 1.718465885608439 * y * z], - [2.864109809347398 * x**3 * w - 1.718465885608439 * x * w], - [2.864109809347398 * y**3 * w - 1.718465885608439 * y * w], - [2.864109809347398 * z**3 * w - 1.718465885608439 * z * w], - [2.864109809347398 * x * w**3 - 1.718465885608439 * x * w], - [2.864109809347398 * y * w**3 - 1.718465885608439 * y * w], - [2.864109809347398 * z * w**3 - 1.718465885608439 * z * w], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - [4.960783708246104 * x**3 * y * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y**3 * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y * z**3 - 2.976470224947662 * x * y * z], - [4.960783708246104 * x**3 * y * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * y**3 * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x**3 * z * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y**3 * z * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * z**3 * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z**3 * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * y * w**3 - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * z * w**3 - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z * w**3 - 2.976470224947662 * y * z * w], - [8.5923294280422 * x**3 * y * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y**3 * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z**3 * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z * w**3 - 5.15539765682532 * x * y * z * w], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624196 * x**2 - 0.2795084971874732], - [0.8385254915624196 * y**2 - 0.2795084971874732], - [0.8385254915624196 * z**2 - 0.2795084971874732], - [0.8385254915624196 * w**2 - 0.2795084971874732], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759272 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759272 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759272 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759272 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759272 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759272 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759272 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759272 * z], - [1.653594569415366 * x**3 - 0.9921567416492196 * x], - [1.653594569415366 * y**3 - 0.9921567416492196 * y], - [1.653594569415366 * z**3 - 0.9921567416492196 * z], - [1.653594569415366 * w**3 - 0.9921567416492196 * w], - [2.25 * x * y * z * w], - [2.515576474687268 * x**2 * y * z - 0.8385254915624226 * y * z], - [2.515576474687268 * x * y**2 * z - 0.8385254915624226 * x * z], - [2.515576474687268 * x * y * z**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x**2 * y * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * x**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * y**2 * z * w - 0.8385254915624226 * z * w], - [2.515576474687268 * x * z**2 * w - 0.8385254915624226 * x * w], - [2.515576474687268 * y * z**2 * w - 0.8385254915624226 * y * w], - [2.515576474687268 * x * y * w**2 - 0.8385254915624226 * x * y], - [2.515576474687268 * x * z * w**2 - 0.8385254915624226 * x * z], - [2.515576474687268 * y * z * w**2 - 0.8385254915624226 * y * z], - [2.8125 * x**2 * y**2 - 0.9375 * y**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * x**2 * z**2 - 0.9375 * z**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * z**2 - 0.9375 * z**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * x**2 * w**2 - 0.9375 * w**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * w**2 - 0.9375 * w**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * z**2 * w**2 - 0.9375 * w**2 - 0.9375 * z**2 + 0.3125], - [2.864109809347398 * x**3 * y - 1.718465885608439 * x * y], - [2.864109809347398 * x * y**3 - 1.718465885608439 * x * y], - [2.864109809347398 * x**3 * z - 1.718465885608439 * x * z], - [2.864109809347398 * y**3 * z - 1.718465885608439 * y * z], - [2.864109809347398 * x * z**3 - 1.718465885608439 * x * z], - [2.864109809347398 * y * z**3 - 1.718465885608439 * y * z], - [2.864109809347398 * x**3 * w - 1.718465885608439 * x * w], - [2.864109809347398 * y**3 * w - 1.718465885608439 * y * w], - [2.864109809347398 * z**3 * w - 1.718465885608439 * z * w], - [2.864109809347398 * x * w**3 - 1.718465885608439 * x * w], - [2.864109809347398 * y * w**3 - 1.718465885608439 * y * w], - [2.864109809347398 * z * w**3 - 1.718465885608439 * z * w], - [3.28125 * x**4 - 2.8125 * x**2 + 0.28125], - [3.28125 * y**4 - 2.8125 * y**2 + 0.28125], - [3.28125 * z**4 - 2.8125 * z**2 + 0.28125], - [3.28125 * w**4 - 2.8125 * w**2 + 0.28125], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - [ - 4.87139289628746 * x**2 * y**2 * z - - 1.62379763209582 * y**2 * z - - 1.62379763209582 * x**2 * z - + 0.5412658773652733 * z - ], - [ - 4.87139289628746 * x**2 * y * z**2 - - 1.62379763209582 * y * z**2 - - 1.62379763209582 * x**2 * y - + 0.5412658773652733 * y - ], - [ - 4.87139289628746 * x * y**2 * z**2 - - 1.62379763209582 * x * z**2 - - 1.62379763209582 * x * y**2 - + 0.5412658773652733 * x - ], - [ - 4.87139289628746 * x**2 * y**2 * w - - 1.62379763209582 * y**2 * w - - 1.62379763209582 * x**2 * w - + 0.5412658773652733 * w - ], - [ - 4.87139289628746 * x**2 * z**2 * w - - 1.62379763209582 * z**2 * w - - 1.62379763209582 * x**2 * w - + 0.5412658773652733 * w - ], - [ - 4.87139289628746 * y**2 * z**2 * w - - 1.62379763209582 * z**2 * w - - 1.62379763209582 * y**2 * w - + 0.5412658773652733 * w - ], - [ - 4.87139289628746 * x**2 * y * w**2 - - 1.62379763209582 * y * w**2 - - 1.62379763209582 * x**2 * y - + 0.5412658773652733 * y - ], - [ - 4.87139289628746 * x * y**2 * w**2 - - 1.62379763209582 * x * w**2 - - 1.62379763209582 * x * y**2 - + 0.5412658773652733 * x - ], - [ - 4.87139289628746 * x**2 * z * w**2 - - 1.62379763209582 * z * w**2 - - 1.62379763209582 * x**2 * z - + 0.5412658773652733 * z - ], - [ - 4.87139289628746 * y**2 * z * w**2 - - 1.62379763209582 * z * w**2 - - 1.62379763209582 * y**2 * z - + 0.5412658773652733 * z - ], - [ - 4.87139289628746 * x * z**2 * w**2 - - 1.62379763209582 * x * w**2 - - 1.62379763209582 * x * z**2 - + 0.5412658773652733 * x - ], - [ - 4.87139289628746 * y * z**2 * w**2 - - 1.62379763209582 * y * w**2 - - 1.62379763209582 * y * z**2 - + 0.5412658773652733 * y - ], - [4.960783708246104 * x**3 * y * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y**3 * z - 2.976470224947662 * x * y * z], - [4.960783708246104 * x * y * z**3 - 2.976470224947662 * x * y * z], - [4.960783708246104 * x**3 * y * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * y**3 * w - 2.976470224947662 * x * y * w], - [4.960783708246104 * x**3 * z * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y**3 * z * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * z**3 * w - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z**3 * w - 2.976470224947662 * y * z * w], - [4.960783708246104 * x * y * w**3 - 2.976470224947662 * x * y * w], - [4.960783708246104 * x * z * w**3 - 2.976470224947662 * x * z * w], - [4.960783708246104 * y * z * w**3 - 2.976470224947662 * y * z * w], - [ - 5.68329171233537 * x**4 * y - - 4.87139289628746 * x**2 * y - + 0.487139289628746 * y - ], - [ - 5.68329171233537 * x * y**4 - - 4.87139289628746 * x * y**2 - + 0.487139289628746 * x - ], - [ - 5.68329171233537 * x**4 * z - - 4.87139289628746 * x**2 * z - + 0.487139289628746 * z - ], - [ - 5.68329171233537 * y**4 * z - - 4.87139289628746 * y**2 * z - + 0.487139289628746 * z - ], - [ - 5.68329171233537 * x * z**4 - - 4.87139289628746 * x * z**2 - + 0.487139289628746 * x - ], - [ - 5.68329171233537 * y * z**4 - - 4.87139289628746 * y * z**2 - + 0.487139289628746 * y - ], - [ - 5.68329171233537 * x**4 * w - - 4.87139289628746 * x**2 * w - + 0.487139289628746 * w - ], - [ - 5.68329171233537 * y**4 * w - - 4.87139289628746 * y**2 * w - + 0.487139289628746 * w - ], - [ - 5.68329171233537 * z**4 * w - - 4.87139289628746 * z**2 * w - + 0.487139289628746 * w - ], - [ - 5.68329171233537 * x * w**4 - - 4.87139289628746 * x * w**2 - + 0.487139289628746 * x - ], - [ - 5.68329171233537 * y * w**4 - - 4.87139289628746 * y * w**2 - + 0.487139289628746 * y - ], - [ - 5.68329171233537 * z * w**4 - - 4.87139289628746 * z * w**2 - + 0.487139289628746 * z - ], - [ - 8.4375 * x**2 * y**2 * z * w - - 2.8125 * y**2 * z * w - - 2.8125 * x**2 * z * w - + 0.9375 * z * w - ], - [ - 8.4375 * x**2 * y * z**2 * w - - 2.8125 * y * z**2 * w - - 2.8125 * x**2 * y * w - + 0.9375 * y * w - ], - [ - 8.4375 * x * y**2 * z**2 * w - - 2.8125 * x * z**2 * w - - 2.8125 * x * y**2 * w - + 0.9375 * x * w - ], - [ - 8.4375 * x**2 * y * z * w**2 - - 2.8125 * y * z * w**2 - - 2.8125 * x**2 * y * z - + 0.9375 * y * z - ], - [ - 8.4375 * x * y**2 * z * w**2 - - 2.8125 * x * z * w**2 - - 2.8125 * x * y**2 * z - + 0.9375 * x * z - ], - [ - 8.4375 * x * y * z**2 * w**2 - - 2.8125 * x * y * w**2 - - 2.8125 * x * y * z**2 - + 0.9375 * x * y - ], - [8.5923294280422 * x**3 * y * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y**3 * z * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z**3 * w - 5.15539765682532 * x * y * z * w], - [8.5923294280422 * x * y * z * w**3 - 5.15539765682532 * x * y * z * w], - [9.84375 * x**4 * y * z - 8.4375 * x**2 * y * z + 0.84375 * y * z], - [9.84375 * x * y**4 * z - 8.4375 * x * y**2 * z + 0.84375 * x * z], - [9.84375 * x * y * z**4 - 8.4375 * x * y * z**2 + 0.84375 * x * y], - [9.84375 * x**4 * y * w - 8.4375 * x**2 * y * w + 0.84375 * y * w], - [9.84375 * x * y**4 * w - 8.4375 * x * y**2 * w + 0.84375 * x * w], - [9.84375 * x**4 * z * w - 8.4375 * x**2 * z * w + 0.84375 * z * w], - [9.84375 * y**4 * z * w - 8.4375 * y**2 * z * w + 0.84375 * z * w], - [9.84375 * x * z**4 * w - 8.4375 * x * z**2 * w + 0.84375 * x * w], - [9.84375 * y * z**4 * w - 8.4375 * y * z**2 * w + 0.84375 * y * w], - [9.84375 * x * y * w**4 - 8.4375 * x * y * w**2 + 0.84375 * x * y], - [9.84375 * x * z * w**4 - 8.4375 * x * z * w**2 + 0.84375 * x * z], - [9.84375 * y * z * w**4 - 8.4375 * y * z * w**2 + 0.84375 * y * z], - [ - 17.04987513700614 * x**4 * y * z * w - - 14.61417868886241 * x**2 * y * z * w - + 1.46141786888624 * y * z * w - ], - [ - 17.04987513700614 * x * y**4 * z * w - - 14.61417868886241 * x * y**2 * z * w - + 1.46141786888624 * x * z * w - ], - [ - 17.04987513700614 * x * y * z**4 * w - - 14.61417868886241 * x * y * z**2 * w - + 1.46141786888624 * x * y * w - ], - [ - 17.04987513700614 * x * y * z * w**4 - - 14.61417868886241 * x * y * z * w**2 - + 1.46141786888624 * x * y * z - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - elif modal and basis_type == "tensor": - if order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922192 * x], - [0.4330127018922192 * y], - [0.4330127018922192 * z], - [0.4330127018922192 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [1.299038105676659 * x * y * z], - [1.299038105676659 * x * y * w], - [1.299038105676659 * x * z * w], - [1.299038105676659 * y * z * w], - [2.25 * x * y * z * w], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922193 * x], - [0.4330127018922193 * y], - [0.4330127018922193 * z], - [0.4330127018922193 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * x * w], - [0.75 * y * w], - [0.75 * z * w], - [0.8385254915624212 * x**2 - 0.2795084971874737], - [0.8385254915624212 * y**2 - 0.2795084971874737], - [0.8385254915624212 * z**2 - 0.2795084971874737], - [0.8385254915624212 * w**2 - 0.2795084971874737], - [1.299038105676658 * x * y * z], - [1.299038105676658 * x * y * w], - [1.299038105676658 * x * z * w], - [1.299038105676658 * y * z * w], - [1.452368754827781 * x**2 * y - 0.4841229182759271 * y], - [1.452368754827781 * x * y**2 - 0.4841229182759271 * x], - [1.452368754827781 * x**2 * z - 0.4841229182759271 * z], - [1.452368754827781 * y**2 * z - 0.4841229182759271 * z], - [1.452368754827781 * x * z**2 - 0.4841229182759271 * x], - [1.452368754827781 * y * z**2 - 0.4841229182759271 * y], - [1.452368754827781 * x**2 * w - 0.4841229182759271 * w], - [1.452368754827781 * y**2 * w - 0.4841229182759271 * w], - [1.452368754827781 * z**2 * w - 0.4841229182759271 * w], - [1.452368754827781 * x * w**2 - 0.4841229182759271 * x], - [1.452368754827781 * y * w**2 - 0.4841229182759271 * y], - [1.452368754827781 * z * w**2 - 0.4841229182759271 * z], - [2.25 * x * y * z * w], - [2.515576474687264 * x**2 * y * z - 0.8385254915624212 * y * z], - [2.515576474687264 * x * y**2 * z - 0.8385254915624212 * x * z], - [2.515576474687264 * x * y * z**2 - 0.8385254915624212 * x * y], - [2.515576474687264 * x**2 * y * w - 0.8385254915624212 * y * w], - [2.515576474687264 * x * y**2 * w - 0.8385254915624212 * x * w], - [2.515576474687264 * x**2 * z * w - 0.8385254915624212 * z * w], - [2.515576474687264 * y**2 * z * w - 0.8385254915624212 * z * w], - [2.515576474687264 * x * z**2 * w - 0.8385254915624212 * x * w], - [2.515576474687264 * y * z**2 * w - 0.8385254915624212 * y * w], - [2.515576474687264 * x * y * w**2 - 0.8385254915624212 * x * y], - [2.515576474687264 * x * z * w**2 - 0.8385254915624212 * x * z], - [2.515576474687264 * y * z * w**2 - 0.8385254915624212 * y * z], - [2.8125 * x**2 * y**2 - 0.9375 * y**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * x**2 * z**2 - 0.9375 * z**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * z**2 - 0.9375 * z**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * x**2 * w**2 - 0.9375 * w**2 - 0.9375 * x**2 + 0.3125], - [2.8125 * y**2 * w**2 - 0.9375 * w**2 - 0.9375 * y**2 + 0.3125], - [2.8125 * z**2 * w**2 - 0.9375 * w**2 - 0.9375 * z**2 + 0.3125], - [4.357106264483344 * x**2 * y * z * w - 1.452368754827781 * y * z * w], - [4.357106264483344 * x * y**2 * z * w - 1.452368754827781 * x * z * w], - [4.357106264483344 * x * y * z**2 * w - 1.452368754827781 * x * y * w], - [4.357106264483344 * x * y * z * w**2 - 1.452368754827781 * x * y * z], - [ - 4.871392896287466 * x**2 * y**2 * z - - 1.623797632095822 * y**2 * z - - 1.623797632095822 * x**2 * z - + 0.541265877365274 * z - ], - [ - 4.871392896287466 * x**2 * y * z**2 - - 1.623797632095822 * y * z**2 - - 1.623797632095822 * x**2 * y - + 0.541265877365274 * y - ], - [ - 4.871392896287466 * x * y**2 * z**2 - - 1.623797632095822 * x * z**2 - - 1.623797632095822 * x * y**2 - + 0.541265877365274 * x - ], - [ - 4.871392896287466 * x**2 * y**2 * w - - 1.623797632095822 * y**2 * w - - 1.623797632095822 * x**2 * w - + 0.541265877365274 * w - ], - [ - 4.871392896287466 * x**2 * z**2 * w - - 1.623797632095822 * z**2 * w - - 1.623797632095822 * x**2 * w - + 0.541265877365274 * w - ], - [ - 4.871392896287466 * y**2 * z**2 * w - - 1.623797632095822 * z**2 * w - - 1.623797632095822 * y**2 * w - + 0.541265877365274 * w - ], - [ - 4.871392896287466 * x**2 * y * w**2 - - 1.623797632095822 * y * w**2 - - 1.623797632095822 * x**2 * y - + 0.541265877365274 * y - ], - [ - 4.871392896287466 * x * y**2 * w**2 - - 1.623797632095822 * x * w**2 - - 1.623797632095822 * x * y**2 - + 0.541265877365274 * x - ], - [ - 4.871392896287466 * x**2 * z * w**2 - - 1.623797632095822 * z * w**2 - - 1.623797632095822 * x**2 * z - + 0.541265877365274 * z - ], - [ - 4.871392896287466 * y**2 * z * w**2 - - 1.623797632095822 * z * w**2 - - 1.623797632095822 * y**2 * z - + 0.541265877365274 * z - ], - [ - 4.871392896287466 * x * z**2 * w**2 - - 1.623797632095822 * x * w**2 - - 1.623797632095822 * x * z**2 - + 0.541265877365274 * x - ], - [ - 4.871392896287466 * y * z**2 * w**2 - - 1.623797632095822 * y * w**2 - - 1.623797632095822 * y * z**2 - + 0.541265877365274 * y - ], - [ - 8.4375 * x**2 * y**2 * z * w - - 2.8125 * y**2 * z * w - - 2.8125 * x**2 * z * w - + 0.9375 * z * w - ], - [ - 8.4375 * x**2 * y * z**2 * w - - 2.8125 * y * z**2 * w - - 2.8125 * x**2 * y * w - + 0.9375 * y * w - ], - [ - 8.4375 * x * y**2 * z**2 * w - - 2.8125 * x * z**2 * w - - 2.8125 * x * y**2 * w - + 0.9375 * x * w - ], - [ - 8.4375 * x**2 * y * z * w**2 - - 2.8125 * y * z * w**2 - - 2.8125 * x**2 * y * z - + 0.9375 * y * z - ], - [ - 8.4375 * x * y**2 * z * w**2 - - 2.8125 * x * z * w**2 - - 2.8125 * x * y**2 * z - + 0.9375 * x * z - ], - [ - 8.4375 * x * y * z**2 * w**2 - - 2.8125 * x * y * w**2 - - 2.8125 * x * y * z**2 - + 0.9375 * x * y - ], - [ - 9.43341178007724 * x**2 * y**2 * z**2 - - 3.14447059335908 * y**2 * z**2 - - 3.14447059335908 * x**2 * z**2 - + 1.048156864453027 * z**2 - - 3.14447059335908 * x**2 * y**2 - + 1.048156864453027 * y**2 - + 1.048156864453027 * x**2 - - 0.3493856214843422 - ], - [ - 9.43341178007724 * x**2 * y**2 * w**2 - - 3.14447059335908 * y**2 * w**2 - - 3.14447059335908 * x**2 * w**2 - + 1.048156864453027 * w**2 - - 3.14447059335908 * x**2 * y**2 - + 1.048156864453027 * y**2 - + 1.048156864453027 * x**2 - - 0.3493856214843422 - ], - [ - 9.43341178007724 * x**2 * z**2 * w**2 - - 3.14447059335908 * z**2 * w**2 - - 3.14447059335908 * x**2 * w**2 - + 1.048156864453027 * w**2 - - 3.14447059335908 * x**2 * z**2 - + 1.048156864453027 * z**2 - + 1.048156864453027 * x**2 - - 0.3493856214843422 - ], - [ - 9.43341178007724 * y**2 * z**2 * w**2 - - 3.14447059335908 * z**2 * w**2 - - 3.14447059335908 * y**2 * w**2 - + 1.048156864453027 * w**2 - - 3.14447059335908 * y**2 * z**2 - + 1.048156864453027 * z**2 - + 1.048156864453027 * y**2 - - 0.3493856214843422 - ], - [ - 16.33914849181254 * x**2 * y**2 * z**2 * w - - 5.44638283060418 * y**2 * z**2 * w - - 5.44638283060418 * x**2 * z**2 * w - + 1.815460943534727 * z**2 * w - - 5.44638283060418 * x**2 * y**2 * w - + 1.815460943534727 * y**2 * w - + 1.815460943534727 * x**2 * w - - 0.6051536478449089 * w - ], - [ - 16.33914849181254 * x**2 * y**2 * z * w**2 - - 5.44638283060418 * y**2 * z * w**2 - - 5.44638283060418 * x**2 * z * w**2 - + 1.815460943534727 * z * w**2 - - 5.44638283060418 * x**2 * y**2 * z - + 1.815460943534727 * y**2 * z - + 1.815460943534727 * x**2 * z - - 0.6051536478449089 * z - ], - [ - 16.33914849181254 * x**2 * y * z**2 * w**2 - - 5.44638283060418 * y * z**2 * w**2 - - 5.44638283060418 * x**2 * y * w**2 - + 1.815460943534727 * y * w**2 - - 5.44638283060418 * x**2 * y * z**2 - + 1.815460943534727 * y * z**2 - + 1.815460943534727 * x**2 * y - - 0.6051536478449089 * y - ], - [ - 16.33914849181254 * x * y**2 * z**2 * w**2 - - 5.44638283060418 * x * z**2 * w**2 - - 5.44638283060418 * x * y**2 * w**2 - + 1.815460943534727 * x * w**2 - - 5.44638283060418 * x * y**2 * z**2 - + 1.815460943534727 * x * z**2 - + 1.815460943534727 * x * y**2 - - 0.6051536478449089 * x - ], - [ - 31.640625 * x**2 * y**2 * z**2 * w**2 - - 10.546875 * y**2 * z**2 * w**2 - - 10.546875 * x**2 * z**2 * w**2 - + 3.515625 * z**2 * w**2 - - 10.546875 * x**2 * y**2 * w**2 - + 3.515625 * y**2 * w**2 - + 3.515625 * x**2 * w**2 - - 1.171875 * w**2 - - 10.546875 * x**2 * y**2 * z**2 - + 3.515625 * y**2 * z**2 - + 3.515625 * x**2 * z**2 - - 1.171875 * z**2 - + 3.515625 * x**2 * y**2 - - 1.171875 * y**2 - - 1.171875 * x**2 - + 0.390625 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <3".format( - order - ) - ) - - elif modal and basis_type == "gkhybrid": - if order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922193 * x], - [0.4330127018922193 * y], - [0.4330127018922193 * z], - [0.4330127018922193 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * w * x], - [0.75 * w * y], - [0.75 * w * z], - [1.299038105676658 * x * y * z], - [1.299038105676658 * w * x * y], - [1.299038105676658 * w * x * z], - [1.299038105676658 * w * y * z], - [2.25 * w * x * y * z], - [0.8385254915624212 * (z**2 - 0.3333333333333333)], - [1.452368754827781 * (x * z**2 - 0.3333333333333333 * x)], - [1.452368754827781 * (y * z**2 - 0.3333333333333333 * y)], - [1.452368754827781 * (w * z**2 - 0.3333333333333333 * w)], - [2.515576474687264 * (x * y * z**2 - 0.3333333333333333 * x * y)], - [2.515576474687264 * (w * x * z**2 - 0.3333333333333333 * w * x)], - [2.515576474687264 * (w * y * z**2 - 0.3333333333333333 * w * y)], - [ - 4.357106264483344 - * (w * x * y * z**2 - 0.3333333333333333 * w * x * y) - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0] - * interpListND[3].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[3].shape[0]): - for j in range(0, interpListND[2].shape[0]): - for k in range(0, interpListND[1].shape[0]): - for l in range(0, interpListND[0].shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpListND[0].shape[0] - + j * interpListND[1].shape[0] * interpListND[0].shape[0] - + i - * interpListND[2].shape[0] - * interpListND[1].shape[0] - * interpListND[0].shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpListND[0][l]) - .subs(y, interpListND[1][k]) - .subs(z, interpListND[2][j]) - .subs(w, interpListND[3][i]) - ) - - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal and basis_type == "hybrid": - if order == 1: - functionVector = Matrix( - [ - [0.25], - [0.4330127018922194 * x], - [0.4330127018922194 * y], - [0.4330127018922194 * z], - [0.4330127018922194 * w], - [0.75 * x * y], - [0.75 * x * z], - [0.75 * y * z], - [0.75 * w * x], - [0.75 * w * y], - [0.75 * w * z], - [1.299038105676658 * x * y * z], - [1.299038105676658 * w * x * y], - [1.299038105676658 * w * x * z], - [1.299038105676658 * w * y * z], - [2.25 * w * x * y * z], - [0.8385254915624211 * (w**2 - 0.3333333333333333)], - [1.452368754827781 * (w**2 * x - 0.3333333333333333 * x)], - [1.452368754827781 * (w**2 * y - 0.3333333333333333 * y)], - [1.452368754827781 * (w**2 * z - 0.3333333333333333 * z)], - [2.515576474687264 * (w**2 * x * y - 0.3333333333333333 * x * y)], - [2.515576474687264 * (w**2 * x * z - 0.3333333333333333 * x * z)], - [2.515576474687264 * (w**2 * y * z - 0.3333333333333333 * y * z)], - [ - 4.357106264483344 - * (w**2 * x * y * z - 0.3333333333333333 * x * y * z) - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0] - * interpListND[3].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[3].shape[0]): - for j in range(0, interpListND[2].shape[0]): - for k in range(0, interpListND[1].shape[0]): - for l in range(0, interpListND[0].shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpListND[0].shape[0] - + j * interpListND[1].shape[0] * interpListND[0].shape[0] - + i - * interpListND[2].shape[0] - * interpListND[1].shape[0] - * interpListND[0].shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpListND[0][l]) - .subs(y, interpListND[1][k]) - .subs(z, interpListND[2][j]) - .subs(w, interpListND[3][i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [ - (w * x) / 16.0 - - x / 16.0 - - y / 16.0 - - z / 16.0 - - w / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - - y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - y / 16.0 - - x / 16.0 - - w / 16.0 - - z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - + y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - z / 16.0 - - x / 16.0 - - y / 16.0 - - w / 16.0 - + (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - - y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - y / 16.0 - - x / 16.0 - - w / 16.0 - + z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - x / 16.0 - - w / 16.0 - + y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - - y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - - y / 16.0 - - z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - + y / 16.0 - - z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - + y / 16.0 - - z / 16.0 - + (w * x) / 16.0 - + (w * y) / 16.0 - - (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - - y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - - (x * z) / 16.0 - - (y * z) / 16.0 - + (w * x * y) / 16.0 - - (w * x * z) / 16.0 - - (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - - y / 16.0 - + z / 16.0 - + (w * x) / 16.0 - - (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - + (x * z) / 16.0 - - (y * z) / 16.0 - - (w * x * y) / 16.0 - + (w * x * z) / 16.0 - - (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - - x / 16.0 - + y / 16.0 - + z / 16.0 - - (w * x) / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - - (x * y) / 16.0 - - (x * z) / 16.0 - + (y * z) / 16.0 - - (w * x * y) / 16.0 - - (w * x * z) / 16.0 - + (w * y * z) / 16.0 - - (x * y * z) / 16.0 - - (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - [ - w / 16.0 - + x / 16.0 - + y / 16.0 - + z / 16.0 - + (w * x) / 16.0 - + (w * y) / 16.0 - + (w * z) / 16.0 - + (x * y) / 16.0 - + (x * z) / 16.0 - + (y * z) / 16.0 - + (w * x * y) / 16.0 - + (w * x * z) / 16.0 - + (w * y * z) / 16.0 - + (x * y * z) / 16.0 - + (w * x * y * z) / 16.0 - + 1.0 / 16.0 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - (w * y) / 8.0 - - y / 8.0 - - z / 8.0 - - w / 8.0 - + (w * z) / 8.0 - + (y * z) / 8.0 - + (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - (w * x) / 8.0 - - x / 8.0 - - z / 8.0 - - w / 8.0 - + (w * z) / 8.0 - + (x * z) / 8.0 - + (w * y**2) / 8.0 - + (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - - z / 8.0 - - (w * x) / 8.0 - + (w * z) / 8.0 - - (x * z) / 8.0 - + (w * y**2) / 8.0 - - (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - y / 8.0 - - w / 8.0 - - z / 8.0 - - (w * y) / 8.0 - + (w * z) / 8.0 - - (y * z) / 8.0 - + (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - - (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - (w * x) / 8.0 - - x / 8.0 - - y / 8.0 - - w / 8.0 - + (w * y) / 8.0 - + (x * y) / 8.0 - + (w * z**2) / 8.0 - + (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - - y / 8.0 - - (w * x) / 8.0 - + (w * y) / 8.0 - - (x * y) / 8.0 - + (w * z**2) / 8.0 - - (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - - w / 8.0 - + (w * x) / 8.0 - - (w * y) / 8.0 - - (x * y) / 8.0 - + (w * z**2) / 8.0 - + (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - + y / 8.0 - - (w * x) / 8.0 - - (w * y) / 8.0 - + (x * y) / 8.0 - + (w * z**2) / 8.0 - - (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - z / 8.0 - - y / 8.0 - - w / 8.0 - + (w * y) / 8.0 - - (w * z) / 8.0 - - (y * z) / 8.0 - + (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - z / 8.0 - - x / 8.0 - - w / 8.0 - + (w * x) / 8.0 - - (w * z) / 8.0 - - (x * z) / 8.0 - + (w * y**2) / 8.0 - + (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - w / 8.0 - + z / 8.0 - - (w * x) / 8.0 - - (w * z) / 8.0 - + (x * z) / 8.0 - + (w * y**2) / 8.0 - - (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - y / 8.0 - - w / 8.0 - + z / 8.0 - - (w * y) / 8.0 - - (w * z) / 8.0 - + (y * z) / 8.0 - + (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - - (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - - (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - - (w * z**2) / 16.0 - + (w * z) / 16.0 - + w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - (x * y) / 8.0 - - y / 8.0 - - z / 8.0 - - x / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - + (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - - z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - - (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - - z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - + (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - - z / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - - (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - + (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - z / 8.0 - - y / 8.0 - - x / 8.0 - + (x * y) / 8.0 - - (x * z) / 8.0 - - (y * z) / 8.0 - + (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - - y / 8.0 - + z / 8.0 - - (x * y) / 8.0 - + (x * z) / 8.0 - - (y * z) / 8.0 - - (w**2 * x) / 8.0 - + (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - + (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - y / 8.0 - - x / 8.0 - + z / 8.0 - - (x * y) / 8.0 - - (x * z) / 8.0 - + (y * z) / 8.0 - + (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - + (w**2 * x * y) / 8.0 - + (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - - (x * y * z) / 8.0 - + (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - x / 8.0 - + y / 8.0 - + z / 8.0 - + (x * y) / 8.0 - + (x * z) / 8.0 - + (y * z) / 8.0 - - (w**2 * x) / 8.0 - - (w**2 * y) / 8.0 - - (w**2 * z) / 8.0 - - w**2 / 8.0 - - (w**2 * x * y) / 8.0 - - (w**2 * x * z) / 8.0 - - (w**2 * y * z) / 8.0 - + (x * y * z) / 8.0 - - (w**2 * x * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - y / 8.0 - - z / 8.0 - - (w * y) / 8.0 - - (w * z) / 8.0 - + (y * z) / 8.0 - - (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - - (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - x / 8.0 - - z / 8.0 - - (w * x) / 8.0 - - (w * z) / 8.0 - + (x * z) / 8.0 - - (w * y**2) / 8.0 - + (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - - z / 8.0 - + (w * x) / 8.0 - - (w * z) / 8.0 - - (x * z) / 8.0 - - (w * y**2) / 8.0 - - (x * y**2) / 8.0 - + (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - + (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - - (x * z) / 16.0 - + x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - + y / 8.0 - - z / 8.0 - + (w * y) / 8.0 - - (w * z) / 8.0 - - (y * z) / 8.0 - - (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - + (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - + (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - - (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - - (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - - (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - + (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - - (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - + (x * z) / 16.0 - - x / 8.0 - - (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - + (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - + z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - x / 8.0 - - y / 8.0 - - (w * x) / 8.0 - - (w * y) / 8.0 - + (x * y) / 8.0 - - (w * z**2) / 8.0 - + (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - - y / 8.0 - + (w * x) / 8.0 - - (w * y) / 8.0 - - (x * y) / 8.0 - - (w * z**2) / 8.0 - - (x * z**2) / 8.0 - + (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - + (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - - x / 8.0 - + y / 8.0 - - (w * x) / 8.0 - + (w * y) / 8.0 - - (x * y) / 8.0 - - (w * z**2) / 8.0 - + (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - + (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - + (x * y * z**2) / 8.0 - - (w * x * y) / 8.0 - + (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - + y / 8.0 - + (w * x) / 8.0 - + (w * y) / 8.0 - + (x * y) / 8.0 - - (w * z**2) / 8.0 - - (x * z**2) / 8.0 - - (y * z**2) / 8.0 - - z**2 / 8.0 - - (w * x * z**2) / 8.0 - - (w * y * z**2) / 8.0 - - (x * y * z**2) / 8.0 - + (w * x * y) / 8.0 - - (w * x * y * z**2) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - y / 8.0 - + z / 8.0 - - (w * y) / 8.0 - + (w * z) / 8.0 - - (y * z) / 8.0 - - (w * x**2) / 8.0 - + (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - + (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - + (x**2 * y * z) / 8.0 - - (w * y * z) / 8.0 - + (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - - (w**2 * y * z) / 16.0 - - (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - - (w * x**2 * y * z) / 16.0 - - (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - - (w * y * z**2) / 16.0 - + (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - - (x**2 * y * z) / 16.0 - - (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - - (y * z**2) / 16.0 - + (y * z) / 16.0 - + y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - - x / 8.0 - + z / 8.0 - - (w * x) / 8.0 - + (w * z) / 8.0 - - (x * z) / 8.0 - - (w * y**2) / 8.0 - + (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - + (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - + (x * y**2 * z) / 8.0 - - (w * x * z) / 8.0 - + (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - w / 8.0 - + x / 8.0 - + z / 8.0 - + (w * x) / 8.0 - + (w * z) / 8.0 - + (x * z) / 8.0 - - (w * y**2) / 8.0 - - (x * y**2) / 8.0 - - (y**2 * z) / 8.0 - - y**2 / 8.0 - - (w * x * y**2) / 8.0 - - (w * y**2 * z) / 8.0 - - (x * y**2 * z) / 8.0 - + (w * x * z) / 8.0 - - (w * x * y**2 * z) / 8.0 - + 1.0 / 8.0 - ], - [ - -(w**2 * x * y * z) / 16.0 - - (w**2 * x * y) / 16.0 - - (w**2 * x * z) / 16.0 - - (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - - (w * x * y**2 * z) / 16.0 - - (w * x * y**2) / 16.0 - - (w * x * y * z**2) / 16.0 - - (w * x * y * z) / 16.0 - - (w * x * z**2) / 16.0 - + (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - - (x * y**2 * z) / 16.0 - - (x * y**2) / 16.0 - - (x * y * z**2) / 16.0 - + (x * y) / 16.0 - - (x * z**2) / 16.0 - + (x * z) / 16.0 - + x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - [ - w / 8.0 - + y / 8.0 - + z / 8.0 - + (w * y) / 8.0 - + (w * z) / 8.0 - + (y * z) / 8.0 - - (w * x**2) / 8.0 - - (x**2 * y) / 8.0 - - (x**2 * z) / 8.0 - - x**2 / 8.0 - - (w * x**2 * y) / 8.0 - - (w * x**2 * z) / 8.0 - - (x**2 * y * z) / 8.0 - + (w * y * z) / 8.0 - - (w * x**2 * y * z) / 8.0 - + 1.0 / 8.0 - ], - [ - (w**2 * x * y * z) / 16.0 - + (w**2 * x * y) / 16.0 - + (w**2 * x * z) / 16.0 - + (w**2 * x) / 16.0 - + (w**2 * y * z) / 16.0 - + (w**2 * y) / 16.0 - + (w**2 * z) / 16.0 - + w**2 / 16.0 - + (w * x**2 * y * z) / 16.0 - + (w * x**2 * y) / 16.0 - + (w * x**2 * z) / 16.0 - + (w * x**2) / 16.0 - + (w * x * y**2 * z) / 16.0 - + (w * x * y**2) / 16.0 - + (w * x * y * z**2) / 16.0 - + (w * x * y * z) / 16.0 - + (w * x * z**2) / 16.0 - - (w * x) / 16.0 - + (w * y**2 * z) / 16.0 - + (w * y**2) / 16.0 - + (w * y * z**2) / 16.0 - - (w * y) / 16.0 - + (w * z**2) / 16.0 - - (w * z) / 16.0 - - w / 8.0 - + (x**2 * y * z) / 16.0 - + (x**2 * y) / 16.0 - + (x**2 * z) / 16.0 - + x**2 / 16.0 - + (x * y**2 * z) / 16.0 - + (x * y**2) / 16.0 - + (x * y * z**2) / 16.0 - - (x * y) / 16.0 - + (x * z**2) / 16.0 - - (x * z) / 16.0 - - x / 8.0 - + (y**2 * z) / 16.0 - + y**2 / 16.0 - + (y * z**2) / 16.0 - - (y * z) / 16.0 - - y / 8.0 - + z**2 / 16.0 - - z / 8.0 - - 3.0 / 16.0 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, functionVector.shape[0]): - interpMatrix[ - l - + k * interpList.shape[0] - + j * interpList.shape[0] * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - m, - ] = ( - functionVector[m] - .subs(x, interpList[l]) - .subs(y, interpList[k]) - .subs(z, interpList[j]) - .subs(w, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <3 for nodal Serendipity in 4D".format( - order - ) - ) - - else: - raise NameError( - "interpMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - - elif dim == 5: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - w = Symbol("w") - v = Symbol("v") - if modal and basis_type == "maximal-order": - if order == 1: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [ - 1.988737822087165 * x**2 * y**2 - - 0.6629126073623886 * y**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * w**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * w**2 - + 0.2209708691207962 - ], - [2.025231468252455 * x**3 * y - 1.215138880951473 * x * y], - [2.025231468252455 * x * y**3 - 1.215138880951473 * x * y], - [2.025231468252455 * x**3 * z - 1.215138880951473 * x * z], - [2.025231468252455 * y**3 * z - 1.215138880951473 * y * z], - [2.025231468252455 * x * z**3 - 1.215138880951473 * x * z], - [2.025231468252455 * y * z**3 - 1.215138880951473 * y * z], - [2.025231468252455 * x**3 * w - 1.215138880951473 * x * w], - [2.025231468252455 * y**3 * w - 1.215138880951473 * y * w], - [2.025231468252455 * z**3 * w - 1.215138880951473 * z * w], - [2.025231468252455 * x * w**3 - 1.215138880951473 * x * w], - [2.025231468252455 * y * w**3 - 1.215138880951473 * y * w], - [2.025231468252455 * z * w**3 - 1.215138880951473 * z * w], - [2.025231468252455 * x**3 * v - 1.215138880951473 * x * v], - [2.025231468252455 * y**3 * v - 1.215138880951473 * y * v], - [2.025231468252455 * z**3 * v - 1.215138880951473 * z * v], - [2.025231468252455 * w**3 * v - 1.215138880951473 * w * v], - [2.025231468252455 * x * v**3 - 1.215138880951473 * x * v], - [2.025231468252455 * y * v**3 - 1.215138880951473 * y * v], - [2.025231468252455 * z * v**3 - 1.215138880951473 * z * v], - [2.025231468252455 * w * v**3 - 1.215138880951473 * w * v], - [ - 2.320194125768356 * x**4 - - 1.988737822087163 * x**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * y**4 - - 1.988737822087163 * y**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * z**4 - - 1.988737822087163 * z**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * w**4 - - 1.988737822087163 * w**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * v**4 - - 1.988737822087163 * v**2 - + 0.1988737822087163 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "serendipity": - if order == 0: - functionVector = Matrix([[0.1767766952966367]]) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 1: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [2.755675960631069 * x * y * z * w * v], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 2: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [2.755675960631069 * x * y * z * w * v], - [3.080939385966559 * x**2 * y * z * w - 1.026979795322186 * y * z * w], - [3.080939385966559 * x * y**2 * z * w - 1.026979795322186 * x * z * w], - [3.080939385966559 * x * y * z**2 * w - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * y * z * w**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x**2 * y * z * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y**2 * z * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * x * y * z**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x**2 * y * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * x**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * y**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * x * z**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * y * z**2 * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y * w**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x * z * w**2 * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * y * z * w**2 * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y * z * v**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x * y * w * v**2 - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * z * w * v**2 - 1.026979795322186 * x * z * w], - [3.080939385966559 * y * z * w * v**2 - 1.026979795322186 * y * z * w], - [ - 5.336343551534144 * x**2 * y * z * w * v - - 1.778781183844715 * y * z * w * v - ], - [ - 5.336343551534144 * x * y**2 * z * w * v - - 1.778781183844715 * x * z * w * v - ], - [ - 5.336343551534144 * x * y * z**2 * w * v - - 1.778781183844715 * x * y * w * v - ], - [ - 5.336343551534144 * x * y * z * w**2 * v - - 1.778781183844715 * x * y * z * v - ], - [ - 5.336343551534144 * x * y * z * w * v**2 - - 1.778781183844715 * x * y * z * w - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 3: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [2.025231468252455 * x**3 * y - 1.215138880951473 * x * y], - [2.025231468252455 * x * y**3 - 1.215138880951473 * x * y], - [2.025231468252455 * x**3 * z - 1.215138880951473 * x * z], - [2.025231468252455 * y**3 * z - 1.215138880951473 * y * z], - [2.025231468252455 * x * z**3 - 1.215138880951473 * x * z], - [2.025231468252455 * y * z**3 - 1.215138880951473 * y * z], - [2.025231468252455 * x**3 * w - 1.215138880951473 * x * w], - [2.025231468252455 * y**3 * w - 1.215138880951473 * y * w], - [2.025231468252455 * z**3 * w - 1.215138880951473 * z * w], - [2.025231468252455 * x * w**3 - 1.215138880951473 * x * w], - [2.025231468252455 * y * w**3 - 1.215138880951473 * y * w], - [2.025231468252455 * z * w**3 - 1.215138880951473 * z * w], - [2.025231468252455 * x**3 * v - 1.215138880951473 * x * v], - [2.025231468252455 * y**3 * v - 1.215138880951473 * y * v], - [2.025231468252455 * z**3 * v - 1.215138880951473 * z * v], - [2.025231468252455 * w**3 * v - 1.215138880951473 * w * v], - [2.025231468252455 * x * v**3 - 1.215138880951473 * x * v], - [2.025231468252455 * y * v**3 - 1.215138880951473 * y * v], - [2.025231468252455 * z * v**3 - 1.215138880951473 * z * v], - [2.025231468252455 * w * v**3 - 1.215138880951473 * w * v], - [2.755675960631069 * x * y * z * w * v], - [3.080939385966559 * x**2 * y * z * w - 1.026979795322186 * y * z * w], - [3.080939385966559 * x * y**2 * z * w - 1.026979795322186 * x * z * w], - [3.080939385966559 * x * y * z**2 * w - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * y * z * w**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x**2 * y * z * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y**2 * z * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * x * y * z**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x**2 * y * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * x**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * y**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * x * z**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * y * z**2 * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y * w**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x * z * w**2 * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * y * z * w**2 * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y * z * v**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x * y * w * v**2 - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * z * w * v**2 - 1.026979795322186 * x * z * w], - [3.080939385966559 * y * z * w * v**2 - 1.026979795322186 * y * z * w], - [3.507803800100568 * x**3 * y * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y**3 * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y * z**3 - 2.104682280060341 * x * y * z], - [3.507803800100568 * x**3 * y * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * y**3 * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x**3 * z * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y**3 * z * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * z**3 * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z**3 * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * y * w**3 - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * z * w**3 - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z * w**3 - 2.104682280060341 * y * z * w], - [3.507803800100568 * x**3 * y * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * y**3 * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x**3 * z * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y**3 * z * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * z**3 * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z**3 * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x**3 * w * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y**3 * w * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z**3 * w * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * w**3 * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w**3 * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w**3 * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * y * v**3 - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * z * v**3 - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z * v**3 - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * w * v**3 - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w * v**3 - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w * v**3 - 2.104682280060341 * z * w * v], - [ - 5.336343551534144 * x**2 * y * z * w * v - - 1.778781183844715 * y * z * w * v - ], - [ - 5.336343551534144 * x * y**2 * z * w * v - - 1.778781183844715 * x * z * w * v - ], - [ - 5.336343551534144 * x * y * z**2 * w * v - - 1.778781183844715 * x * y * w * v - ], - [ - 5.336343551534144 * x * y * z * w**2 * v - - 1.778781183844715 * x * y * z * v - ], - [ - 5.336343551534144 * x * y * z * w * v**2 - - 1.778781183844715 * x * y * z * w - ], - [ - 6.075694404757367 * x**3 * y * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y**3 * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z**3 * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z * w**3 - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x**3 * y * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y**3 * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * z**3 * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x**3 * y * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * y**3 * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x**3 * z * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y**3 * z * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * z**3 * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z**3 * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * w**3 * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w**3 * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w**3 * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * z * v**3 - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * w * v**3 - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w * v**3 - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w * v**3 - - 3.64541664285442 * y * z * w * v - ], - [ - 10.52341140030171 * x**3 * y * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y**3 * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z**3 * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w**3 * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w * v**3 - - 6.314046840181025 * x * y * z * w * v - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - - elif order == 4: - functionVector = Matrix( - [ - [0.1767766952966367], - [0.3061862178478966 * x], - [0.3061862178478966 * y], - [0.3061862178478966 * z], - [0.3061862178478966 * w], - [0.3061862178478966 * v], - [0.5303300858899102 * x * y], - [0.5303300858899102 * x * z], - [0.5303300858899102 * y * z], - [0.5303300858899102 * x * w], - [0.5303300858899102 * y * w], - [0.5303300858899102 * z * w], - [0.5303300858899102 * x * v], - [0.5303300858899102 * y * v], - [0.5303300858899102 * z * v], - [0.5303300858899102 * w * v], - [0.592927061281571 * x**2 - 0.1976423537605237], - [0.592927061281571 * y**2 - 0.1976423537605237], - [0.592927061281571 * z**2 - 0.1976423537605237], - [0.592927061281571 * w**2 - 0.1976423537605237], - [0.592927061281571 * v**2 - 0.1976423537605237], - [0.9185586535436896 * x * y * z], - [0.9185586535436896 * x * y * w], - [0.9185586535436896 * x * z * w], - [0.9185586535436896 * y * z * w], - [0.9185586535436896 * x * y * v], - [0.9185586535436896 * x * z * v], - [0.9185586535436896 * y * z * v], - [0.9185586535436896 * x * w * v], - [0.9185586535436896 * y * w * v], - [0.9185586535436896 * z * w * v], - [1.026979795322187 * x**2 * y - 0.3423265984407291 * y], - [1.026979795322187 * x * y**2 - 0.3423265984407291 * x], - [1.026979795322187 * x**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * y**2 * z - 0.3423265984407291 * z], - [1.026979795322187 * x * z**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * z**2 - 0.3423265984407291 * y], - [1.026979795322187 * x**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * y**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * z**2 * w - 0.3423265984407291 * w], - [1.026979795322187 * x * w**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * w**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * w**2 - 0.3423265984407291 * z], - [1.026979795322187 * x**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * y**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * z**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * w**2 * v - 0.3423265984407291 * v], - [1.026979795322187 * x * v**2 - 0.3423265984407291 * x], - [1.026979795322187 * y * v**2 - 0.3423265984407291 * y], - [1.026979795322187 * z * v**2 - 0.3423265984407291 * z], - [1.026979795322187 * w * v**2 - 0.3423265984407291 * w], - [1.169267933366857 * x**3 - 0.701560760020114 * x], - [1.169267933366857 * y**3 - 0.701560760020114 * y], - [1.169267933366857 * z**3 - 0.701560760020114 * z], - [1.169267933366857 * w**3 - 0.701560760020114 * w], - [1.169267933366857 * v**3 - 0.701560760020114 * v], - [1.590990257669732 * x * y * z * w], - [1.590990257669732 * x * y * z * v], - [1.590990257669732 * x * y * w * v], - [1.590990257669732 * x * z * w * v], - [1.590990257669732 * y * z * w * v], - [1.778781183844712 * x**2 * y * z - 0.5929270612815707 * y * z], - [1.778781183844712 * x * y**2 * z - 0.5929270612815707 * x * z], - [1.778781183844712 * x * y * z**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x**2 * y * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * x**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * y**2 * z * w - 0.5929270612815707 * z * w], - [1.778781183844712 * x * z**2 * w - 0.5929270612815707 * x * w], - [1.778781183844712 * y * z**2 * w - 0.5929270612815707 * y * w], - [1.778781183844712 * x * y * w**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * w**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * w**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x**2 * y * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x * y**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * x**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * y**2 * z * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * z**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * z**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * x**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * y**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * z**2 * w * v - 0.5929270612815707 * w * v], - [1.778781183844712 * x * w**2 * v - 0.5929270612815707 * x * v], - [1.778781183844712 * y * w**2 * v - 0.5929270612815707 * y * v], - [1.778781183844712 * z * w**2 * v - 0.5929270612815707 * z * v], - [1.778781183844712 * x * y * v**2 - 0.5929270612815707 * x * y], - [1.778781183844712 * x * z * v**2 - 0.5929270612815707 * x * z], - [1.778781183844712 * y * z * v**2 - 0.5929270612815707 * y * z], - [1.778781183844712 * x * w * v**2 - 0.5929270612815707 * x * w], - [1.778781183844712 * y * w * v**2 - 0.5929270612815707 * y * w], - [1.778781183844712 * z * w * v**2 - 0.5929270612815707 * z * w], - [ - 1.988737822087165 * x**2 * y**2 - - 0.6629126073623886 * y**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * z**2 - - 0.6629126073623886 * z**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * w**2 - - 0.6629126073623886 * w**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * x**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * x**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * y**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * y**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * z**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * z**2 - + 0.2209708691207962 - ], - [ - 1.988737822087165 * w**2 * v**2 - - 0.6629126073623886 * v**2 - - 0.6629126073623886 * w**2 - + 0.2209708691207962 - ], - [2.025231468252455 * x**3 * y - 1.215138880951473 * x * y], - [2.025231468252455 * x * y**3 - 1.215138880951473 * x * y], - [2.025231468252455 * x**3 * z - 1.215138880951473 * x * z], - [2.025231468252455 * y**3 * z - 1.215138880951473 * y * z], - [2.025231468252455 * x * z**3 - 1.215138880951473 * x * z], - [2.025231468252455 * y * z**3 - 1.215138880951473 * y * z], - [2.025231468252455 * x**3 * w - 1.215138880951473 * x * w], - [2.025231468252455 * y**3 * w - 1.215138880951473 * y * w], - [2.025231468252455 * z**3 * w - 1.215138880951473 * z * w], - [2.025231468252455 * x * w**3 - 1.215138880951473 * x * w], - [2.025231468252455 * y * w**3 - 1.215138880951473 * y * w], - [2.025231468252455 * z * w**3 - 1.215138880951473 * z * w], - [2.025231468252455 * x**3 * v - 1.215138880951473 * x * v], - [2.025231468252455 * y**3 * v - 1.215138880951473 * y * v], - [2.025231468252455 * z**3 * v - 1.215138880951473 * z * v], - [2.025231468252455 * w**3 * v - 1.215138880951473 * w * v], - [2.025231468252455 * x * v**3 - 1.215138880951473 * x * v], - [2.025231468252455 * y * v**3 - 1.215138880951473 * y * v], - [2.025231468252455 * z * v**3 - 1.215138880951473 * z * v], - [2.025231468252455 * w * v**3 - 1.215138880951473 * w * v], - [ - 2.320194125768356 * x**4 - - 1.988737822087163 * x**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * y**4 - - 1.988737822087163 * y**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * z**4 - - 1.988737822087163 * z**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * w**4 - - 1.988737822087163 * w**2 - + 0.1988737822087163 - ], - [ - 2.320194125768356 * v**4 - - 1.988737822087163 * v**2 - + 0.1988737822087163 - ], - [2.755675960631069 * x * y * z * w * v], - [3.080939385966559 * x**2 * y * z * w - 1.026979795322186 * y * z * w], - [3.080939385966559 * x * y**2 * z * w - 1.026979795322186 * x * z * w], - [3.080939385966559 * x * y * z**2 * w - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * y * z * w**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x**2 * y * z * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y**2 * z * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * x * y * z**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x**2 * y * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * x**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * y**2 * z * w * v - 1.026979795322186 * z * w * v], - [3.080939385966559 * x * z**2 * w * v - 1.026979795322186 * x * w * v], - [3.080939385966559 * y * z**2 * w * v - 1.026979795322186 * y * w * v], - [3.080939385966559 * x * y * w**2 * v - 1.026979795322186 * x * y * v], - [3.080939385966559 * x * z * w**2 * v - 1.026979795322186 * x * z * v], - [3.080939385966559 * y * z * w**2 * v - 1.026979795322186 * y * z * v], - [3.080939385966559 * x * y * z * v**2 - 1.026979795322186 * x * y * z], - [3.080939385966559 * x * y * w * v**2 - 1.026979795322186 * x * y * w], - [3.080939385966559 * x * z * w * v**2 - 1.026979795322186 * x * z * w], - [3.080939385966559 * y * z * w * v**2 - 1.026979795322186 * y * z * w], - [ - 3.444594950788842 * x**2 * y**2 * z - - 1.148198316929614 * y**2 * z - - 1.148198316929614 * x**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * x**2 * y * z**2 - - 1.148198316929614 * y * z**2 - - 1.148198316929614 * x**2 * y - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x * y**2 * z**2 - - 1.148198316929614 * x * z**2 - - 1.148198316929614 * x * y**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * x**2 * y**2 * w - - 1.148198316929614 * y**2 * w - - 1.148198316929614 * x**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * x**2 * z**2 * w - - 1.148198316929614 * z**2 * w - - 1.148198316929614 * x**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * y**2 * z**2 * w - - 1.148198316929614 * z**2 * w - - 1.148198316929614 * y**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * x**2 * y * w**2 - - 1.148198316929614 * y * w**2 - - 1.148198316929614 * x**2 * y - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x * y**2 * w**2 - - 1.148198316929614 * x * w**2 - - 1.148198316929614 * x * y**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * x**2 * z * w**2 - - 1.148198316929614 * z * w**2 - - 1.148198316929614 * x**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * y**2 * z * w**2 - - 1.148198316929614 * z * w**2 - - 1.148198316929614 * y**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * x * z**2 * w**2 - - 1.148198316929614 * x * w**2 - - 1.148198316929614 * x * z**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * y * z**2 * w**2 - - 1.148198316929614 * y * w**2 - - 1.148198316929614 * y * z**2 - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x**2 * y**2 * v - - 1.148198316929614 * y**2 * v - - 1.148198316929614 * x**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * x**2 * z**2 * v - - 1.148198316929614 * z**2 * v - - 1.148198316929614 * x**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * y**2 * z**2 * v - - 1.148198316929614 * z**2 * v - - 1.148198316929614 * y**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * x**2 * w**2 * v - - 1.148198316929614 * w**2 * v - - 1.148198316929614 * x**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * y**2 * w**2 * v - - 1.148198316929614 * w**2 * v - - 1.148198316929614 * y**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * z**2 * w**2 * v - - 1.148198316929614 * w**2 * v - - 1.148198316929614 * z**2 * v - + 0.3827327723098713 * v - ], - [ - 3.444594950788842 * x**2 * y * v**2 - - 1.148198316929614 * y * v**2 - - 1.148198316929614 * x**2 * y - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x * y**2 * v**2 - - 1.148198316929614 * x * v**2 - - 1.148198316929614 * x * y**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * x**2 * z * v**2 - - 1.148198316929614 * z * v**2 - - 1.148198316929614 * x**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * y**2 * z * v**2 - - 1.148198316929614 * z * v**2 - - 1.148198316929614 * y**2 * z - + 0.3827327723098713 * z - ], - [ - 3.444594950788842 * x * z**2 * v**2 - - 1.148198316929614 * x * v**2 - - 1.148198316929614 * x * z**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * y * z**2 * v**2 - - 1.148198316929614 * y * v**2 - - 1.148198316929614 * y * z**2 - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * x**2 * w * v**2 - - 1.148198316929614 * w * v**2 - - 1.148198316929614 * x**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * y**2 * w * v**2 - - 1.148198316929614 * w * v**2 - - 1.148198316929614 * y**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * z**2 * w * v**2 - - 1.148198316929614 * w * v**2 - - 1.148198316929614 * z**2 * w - + 0.3827327723098713 * w - ], - [ - 3.444594950788842 * x * w**2 * v**2 - - 1.148198316929614 * x * v**2 - - 1.148198316929614 * x * w**2 - + 0.3827327723098713 * x - ], - [ - 3.444594950788842 * y * w**2 * v**2 - - 1.148198316929614 * y * v**2 - - 1.148198316929614 * y * w**2 - + 0.3827327723098713 * y - ], - [ - 3.444594950788842 * z * w**2 * v**2 - - 1.148198316929614 * z * v**2 - - 1.148198316929614 * z * w**2 - + 0.3827327723098713 * z - ], - [3.507803800100568 * x**3 * y * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y**3 * z - 2.104682280060341 * x * y * z], - [3.507803800100568 * x * y * z**3 - 2.104682280060341 * x * y * z], - [3.507803800100568 * x**3 * y * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * y**3 * w - 2.104682280060341 * x * y * w], - [3.507803800100568 * x**3 * z * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y**3 * z * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * z**3 * w - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z**3 * w - 2.104682280060341 * y * z * w], - [3.507803800100568 * x * y * w**3 - 2.104682280060341 * x * y * w], - [3.507803800100568 * x * z * w**3 - 2.104682280060341 * x * z * w], - [3.507803800100568 * y * z * w**3 - 2.104682280060341 * y * z * w], - [3.507803800100568 * x**3 * y * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * y**3 * v - 2.104682280060341 * x * y * v], - [3.507803800100568 * x**3 * z * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y**3 * z * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * z**3 * v - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z**3 * v - 2.104682280060341 * y * z * v], - [3.507803800100568 * x**3 * w * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y**3 * w * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z**3 * w * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * w**3 * v - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w**3 * v - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w**3 * v - 2.104682280060341 * z * w * v], - [3.507803800100568 * x * y * v**3 - 2.104682280060341 * x * y * v], - [3.507803800100568 * x * z * v**3 - 2.104682280060341 * x * z * v], - [3.507803800100568 * y * z * v**3 - 2.104682280060341 * y * z * v], - [3.507803800100568 * x * w * v**3 - 2.104682280060341 * x * w * v], - [3.507803800100568 * y * w * v**3 - 2.104682280060341 * y * w * v], - [3.507803800100568 * z * w * v**3 - 2.104682280060341 * z * w * v], - [ - 4.018694109253645 * x**4 * y - - 3.444594950788839 * x**2 * y - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * x * y**4 - - 3.444594950788839 * x * y**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * x**4 * z - - 3.444594950788839 * x**2 * z - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * y**4 * z - - 3.444594950788839 * y**2 * z - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * x * z**4 - - 3.444594950788839 * x * z**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * y * z**4 - - 3.444594950788839 * y * z**2 - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * x**4 * w - - 3.444594950788839 * x**2 * w - + 0.3444594950788838 * w - ], - [ - 4.018694109253645 * y**4 * w - - 3.444594950788839 * y**2 * w - + 0.3444594950788838 * w - ], - [ - 4.018694109253645 * z**4 * w - - 3.444594950788839 * z**2 * w - + 0.3444594950788838 * w - ], - [ - 4.018694109253645 * x * w**4 - - 3.444594950788839 * x * w**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * y * w**4 - - 3.444594950788839 * y * w**2 - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * z * w**4 - - 3.444594950788839 * z * w**2 - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * x**4 * v - - 3.444594950788839 * x**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * y**4 * v - - 3.444594950788839 * y**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * z**4 * v - - 3.444594950788839 * z**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * w**4 * v - - 3.444594950788839 * w**2 * v - + 0.3444594950788838 * v - ], - [ - 4.018694109253645 * x * v**4 - - 3.444594950788839 * x * v**2 - + 0.3444594950788838 * x - ], - [ - 4.018694109253645 * y * v**4 - - 3.444594950788839 * y * v**2 - + 0.3444594950788838 * y - ], - [ - 4.018694109253645 * z * v**4 - - 3.444594950788839 * z * v**2 - + 0.3444594950788838 * z - ], - [ - 4.018694109253645 * w * v**4 - - 3.444594950788839 * w * v**2 - + 0.3444594950788838 * w - ], - [ - 5.336343551534144 * x**2 * y * z * w * v - - 1.778781183844715 * y * z * w * v - ], - [ - 5.336343551534144 * x * y**2 * z * w * v - - 1.778781183844715 * x * z * w * v - ], - [ - 5.336343551534144 * x * y * z**2 * w * v - - 1.778781183844715 * x * y * w * v - ], - [ - 5.336343551534144 * x * y * z * w**2 * v - - 1.778781183844715 * x * y * z * v - ], - [ - 5.336343551534144 * x * y * z * w * v**2 - - 1.778781183844715 * x * y * z * w - ], - [ - 5.966213466261497 * x**2 * y**2 * z * w - - 1.988737822087165 * y**2 * z * w - - 1.988737822087165 * x**2 * z * w - + 0.6629126073623886 * z * w - ], - [ - 5.966213466261497 * x**2 * y * z**2 * w - - 1.988737822087165 * y * z**2 * w - - 1.988737822087165 * x**2 * y * w - + 0.6629126073623886 * y * w - ], - [ - 5.966213466261497 * x * y**2 * z**2 * w - - 1.988737822087165 * x * z**2 * w - - 1.988737822087165 * x * y**2 * w - + 0.6629126073623886 * x * w - ], - [ - 5.966213466261497 * x**2 * y * z * w**2 - - 1.988737822087165 * y * z * w**2 - - 1.988737822087165 * x**2 * y * z - + 0.6629126073623886 * y * z - ], - [ - 5.966213466261497 * x * y**2 * z * w**2 - - 1.988737822087165 * x * z * w**2 - - 1.988737822087165 * x * y**2 * z - + 0.6629126073623886 * x * z - ], - [ - 5.966213466261497 * x * y * z**2 * w**2 - - 1.988737822087165 * x * y * w**2 - - 1.988737822087165 * x * y * z**2 - + 0.6629126073623886 * x * y - ], - [ - 5.966213466261497 * x**2 * y**2 * z * v - - 1.988737822087165 * y**2 * z * v - - 1.988737822087165 * x**2 * z * v - + 0.6629126073623886 * z * v - ], - [ - 5.966213466261497 * x**2 * y * z**2 * v - - 1.988737822087165 * y * z**2 * v - - 1.988737822087165 * x**2 * y * v - + 0.6629126073623886 * y * v - ], - [ - 5.966213466261497 * x * y**2 * z**2 * v - - 1.988737822087165 * x * z**2 * v - - 1.988737822087165 * x * y**2 * v - + 0.6629126073623886 * x * v - ], - [ - 5.966213466261497 * x**2 * y**2 * w * v - - 1.988737822087165 * y**2 * w * v - - 1.988737822087165 * x**2 * w * v - + 0.6629126073623886 * w * v - ], - [ - 5.966213466261497 * x**2 * z**2 * w * v - - 1.988737822087165 * z**2 * w * v - - 1.988737822087165 * x**2 * w * v - + 0.6629126073623886 * w * v - ], - [ - 5.966213466261497 * y**2 * z**2 * w * v - - 1.988737822087165 * z**2 * w * v - - 1.988737822087165 * y**2 * w * v - + 0.6629126073623886 * w * v - ], - [ - 5.966213466261497 * x**2 * y * w**2 * v - - 1.988737822087165 * y * w**2 * v - - 1.988737822087165 * x**2 * y * v - + 0.6629126073623886 * y * v - ], - [ - 5.966213466261497 * x * y**2 * w**2 * v - - 1.988737822087165 * x * w**2 * v - - 1.988737822087165 * x * y**2 * v - + 0.6629126073623886 * x * v - ], - [ - 5.966213466261497 * x**2 * z * w**2 * v - - 1.988737822087165 * z * w**2 * v - - 1.988737822087165 * x**2 * z * v - + 0.6629126073623886 * z * v - ], - [ - 5.966213466261497 * y**2 * z * w**2 * v - - 1.988737822087165 * z * w**2 * v - - 1.988737822087165 * y**2 * z * v - + 0.6629126073623886 * z * v - ], - [ - 5.966213466261497 * x * z**2 * w**2 * v - - 1.988737822087165 * x * w**2 * v - - 1.988737822087165 * x * z**2 * v - + 0.6629126073623886 * x * v - ], - [ - 5.966213466261497 * y * z**2 * w**2 * v - - 1.988737822087165 * y * w**2 * v - - 1.988737822087165 * y * z**2 * v - + 0.6629126073623886 * y * v - ], - [ - 5.966213466261497 * x**2 * y * z * v**2 - - 1.988737822087165 * y * z * v**2 - - 1.988737822087165 * x**2 * y * z - + 0.6629126073623886 * y * z - ], - [ - 5.966213466261497 * x * y**2 * z * v**2 - - 1.988737822087165 * x * z * v**2 - - 1.988737822087165 * x * y**2 * z - + 0.6629126073623886 * x * z - ], - [ - 5.966213466261497 * x * y * z**2 * v**2 - - 1.988737822087165 * x * y * v**2 - - 1.988737822087165 * x * y * z**2 - + 0.6629126073623886 * x * y - ], - [ - 5.966213466261497 * x**2 * y * w * v**2 - - 1.988737822087165 * y * w * v**2 - - 1.988737822087165 * x**2 * y * w - + 0.6629126073623886 * y * w - ], - [ - 5.966213466261497 * x * y**2 * w * v**2 - - 1.988737822087165 * x * w * v**2 - - 1.988737822087165 * x * y**2 * w - + 0.6629126073623886 * x * w - ], - [ - 5.966213466261497 * x**2 * z * w * v**2 - - 1.988737822087165 * z * w * v**2 - - 1.988737822087165 * x**2 * z * w - + 0.6629126073623886 * z * w - ], - [ - 5.966213466261497 * y**2 * z * w * v**2 - - 1.988737822087165 * z * w * v**2 - - 1.988737822087165 * y**2 * z * w - + 0.6629126073623886 * z * w - ], - [ - 5.966213466261497 * x * z**2 * w * v**2 - - 1.988737822087165 * x * w * v**2 - - 1.988737822087165 * x * z**2 * w - + 0.6629126073623886 * x * w - ], - [ - 5.966213466261497 * y * z**2 * w * v**2 - - 1.988737822087165 * y * w * v**2 - - 1.988737822087165 * y * z**2 * w - + 0.6629126073623886 * y * w - ], - [ - 5.966213466261497 * x * y * w**2 * v**2 - - 1.988737822087165 * x * y * v**2 - - 1.988737822087165 * x * y * w**2 - + 0.6629126073623886 * x * y - ], - [ - 5.966213466261497 * x * z * w**2 * v**2 - - 1.988737822087165 * x * z * v**2 - - 1.988737822087165 * x * z * w**2 - + 0.6629126073623886 * x * z - ], - [ - 5.966213466261497 * y * z * w**2 * v**2 - - 1.988737822087165 * y * z * v**2 - - 1.988737822087165 * y * z * w**2 - + 0.6629126073623886 * y * z - ], - [ - 6.075694404757367 * x**3 * y * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y**3 * z * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z**3 * w - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x * y * z * w**3 - - 3.64541664285442 * x * y * z * w - ], - [ - 6.075694404757367 * x**3 * y * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y**3 * z * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * z**3 * v - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x**3 * y * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * y**3 * w * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x**3 * z * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y**3 * z * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * z**3 * w * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z**3 * w * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * w**3 * v - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w**3 * v - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w**3 * v - - 3.64541664285442 * y * z * w * v - ], - [ - 6.075694404757367 * x * y * z * v**3 - - 3.64541664285442 * x * y * z * v - ], - [ - 6.075694404757367 * x * y * w * v**3 - - 3.64541664285442 * x * y * w * v - ], - [ - 6.075694404757367 * x * z * w * v**3 - - 3.64541664285442 * x * z * w * v - ], - [ - 6.075694404757367 * y * z * w * v**3 - - 3.64541664285442 * y * z * w * v - ], - [ - 6.960582377305069 * x**4 * y * z - - 5.966213466261488 * x**2 * y * z - + 0.5966213466261489 * y * z - ], - [ - 6.960582377305069 * x * y**4 * z - - 5.966213466261488 * x * y**2 * z - + 0.5966213466261489 * x * z - ], - [ - 6.960582377305069 * x * y * z**4 - - 5.966213466261488 * x * y * z**2 - + 0.5966213466261489 * x * y - ], - [ - 6.960582377305069 * x**4 * y * w - - 5.966213466261488 * x**2 * y * w - + 0.5966213466261489 * y * w - ], - [ - 6.960582377305069 * x * y**4 * w - - 5.966213466261488 * x * y**2 * w - + 0.5966213466261489 * x * w - ], - [ - 6.960582377305069 * x**4 * z * w - - 5.966213466261488 * x**2 * z * w - + 0.5966213466261489 * z * w - ], - [ - 6.960582377305069 * y**4 * z * w - - 5.966213466261488 * y**2 * z * w - + 0.5966213466261489 * z * w - ], - [ - 6.960582377305069 * x * z**4 * w - - 5.966213466261488 * x * z**2 * w - + 0.5966213466261489 * x * w - ], - [ - 6.960582377305069 * y * z**4 * w - - 5.966213466261488 * y * z**2 * w - + 0.5966213466261489 * y * w - ], - [ - 6.960582377305069 * x * y * w**4 - - 5.966213466261488 * x * y * w**2 - + 0.5966213466261489 * x * y - ], - [ - 6.960582377305069 * x * z * w**4 - - 5.966213466261488 * x * z * w**2 - + 0.5966213466261489 * x * z - ], - [ - 6.960582377305069 * y * z * w**4 - - 5.966213466261488 * y * z * w**2 - + 0.5966213466261489 * y * z - ], - [ - 6.960582377305069 * x**4 * y * v - - 5.966213466261488 * x**2 * y * v - + 0.5966213466261489 * y * v - ], - [ - 6.960582377305069 * x * y**4 * v - - 5.966213466261488 * x * y**2 * v - + 0.5966213466261489 * x * v - ], - [ - 6.960582377305069 * x**4 * z * v - - 5.966213466261488 * x**2 * z * v - + 0.5966213466261489 * z * v - ], - [ - 6.960582377305069 * y**4 * z * v - - 5.966213466261488 * y**2 * z * v - + 0.5966213466261489 * z * v - ], - [ - 6.960582377305069 * x * z**4 * v - - 5.966213466261488 * x * z**2 * v - + 0.5966213466261489 * x * v - ], - [ - 6.960582377305069 * y * z**4 * v - - 5.966213466261488 * y * z**2 * v - + 0.5966213466261489 * y * v - ], - [ - 6.960582377305069 * x**4 * w * v - - 5.966213466261488 * x**2 * w * v - + 0.5966213466261489 * w * v - ], - [ - 6.960582377305069 * y**4 * w * v - - 5.966213466261488 * y**2 * w * v - + 0.5966213466261489 * w * v - ], - [ - 6.960582377305069 * z**4 * w * v - - 5.966213466261488 * z**2 * w * v - + 0.5966213466261489 * w * v - ], - [ - 6.960582377305069 * x * w**4 * v - - 5.966213466261488 * x * w**2 * v - + 0.5966213466261489 * x * v - ], - [ - 6.960582377305069 * y * w**4 * v - - 5.966213466261488 * y * w**2 * v - + 0.5966213466261489 * y * v - ], - [ - 6.960582377305069 * z * w**4 * v - - 5.966213466261488 * z * w**2 * v - + 0.5966213466261489 * z * v - ], - [ - 6.960582377305069 * x * y * v**4 - - 5.966213466261488 * x * y * v**2 - + 0.5966213466261489 * x * y - ], - [ - 6.960582377305069 * x * z * v**4 - - 5.966213466261488 * x * z * v**2 - + 0.5966213466261489 * x * z - ], - [ - 6.960582377305069 * y * z * v**4 - - 5.966213466261488 * y * z * v**2 - + 0.5966213466261489 * y * z - ], - [ - 6.960582377305069 * x * w * v**4 - - 5.966213466261488 * x * w * v**2 - + 0.5966213466261489 * x * w - ], - [ - 6.960582377305069 * y * w * v**4 - - 5.966213466261488 * y * w * v**2 - + 0.5966213466261489 * y * w - ], - [ - 6.960582377305069 * z * w * v**4 - - 5.966213466261488 * z * w * v**2 - + 0.5966213466261489 * z * w - ], - [ - 10.33378485236653 * x**2 * y**2 * z * w * v - - 3.444594950788842 * y**2 * z * w * v - - 3.444594950788842 * x**2 * z * w * v - + 1.148198316929614 * z * w * v - ], - [ - 10.33378485236653 * x**2 * y * z**2 * w * v - - 3.444594950788842 * y * z**2 * w * v - - 3.444594950788842 * x**2 * y * w * v - + 1.148198316929614 * y * w * v - ], - [ - 10.33378485236653 * x * y**2 * z**2 * w * v - - 3.444594950788842 * x * z**2 * w * v - - 3.444594950788842 * x * y**2 * w * v - + 1.148198316929614 * x * w * v - ], - [ - 10.33378485236653 * x**2 * y * z * w**2 * v - - 3.444594950788842 * y * z * w**2 * v - - 3.444594950788842 * x**2 * y * z * v - + 1.148198316929614 * y * z * v - ], - [ - 10.33378485236653 * x * y**2 * z * w**2 * v - - 3.444594950788842 * x * z * w**2 * v - - 3.444594950788842 * x * y**2 * z * v - + 1.148198316929614 * x * z * v - ], - [ - 10.33378485236653 * x * y * z**2 * w**2 * v - - 3.444594950788842 * x * y * w**2 * v - - 3.444594950788842 * x * y * z**2 * v - + 1.148198316929614 * x * y * v - ], - [ - 10.33378485236653 * x**2 * y * z * w * v**2 - - 3.444594950788842 * y * z * w * v**2 - - 3.444594950788842 * x**2 * y * z * w - + 1.148198316929614 * y * z * w - ], - [ - 10.33378485236653 * x * y**2 * z * w * v**2 - - 3.444594950788842 * x * z * w * v**2 - - 3.444594950788842 * x * y**2 * z * w - + 1.148198316929614 * x * z * w - ], - [ - 10.33378485236653 * x * y * z**2 * w * v**2 - - 3.444594950788842 * x * y * w * v**2 - - 3.444594950788842 * x * y * z**2 * w - + 1.148198316929614 * x * y * w - ], - [ - 10.33378485236653 * x * y * z * w**2 * v**2 - - 3.444594950788842 * x * y * z * v**2 - - 3.444594950788842 * x * y * z * w**2 - + 1.148198316929614 * x * y * z - ], - [ - 10.52341140030171 * x**3 * y * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y**3 * z * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z**3 * w * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w**3 * v - - 6.314046840181025 * x * y * z * w * v - ], - [ - 10.52341140030171 * x * y * z * w * v**3 - - 6.314046840181025 * x * y * z * w * v - ], - [ - 12.05608232776096 * x**4 * y * z * w - - 10.33378485236654 * x**2 * y * z * w - + 1.033378485236654 * y * z * w - ], - [ - 12.05608232776096 * x * y**4 * z * w - - 10.33378485236654 * x * y**2 * z * w - + 1.033378485236654 * x * z * w - ], - [ - 12.05608232776096 * x * y * z**4 * w - - 10.33378485236654 * x * y * z**2 * w - + 1.033378485236654 * x * y * w - ], - [ - 12.05608232776096 * x * y * z * w**4 - - 10.33378485236654 * x * y * z * w**2 - + 1.033378485236654 * x * y * z - ], - [ - 12.05608232776096 * x**4 * y * z * v - - 10.33378485236654 * x**2 * y * z * v - + 1.033378485236654 * y * z * v - ], - [ - 12.05608232776096 * x * y**4 * z * v - - 10.33378485236654 * x * y**2 * z * v - + 1.033378485236654 * x * z * v - ], - [ - 12.05608232776096 * x * y * z**4 * v - - 10.33378485236654 * x * y * z**2 * v - + 1.033378485236654 * x * y * v - ], - [ - 12.05608232776096 * x**4 * y * w * v - - 10.33378485236654 * x**2 * y * w * v - + 1.033378485236654 * y * w * v - ], - [ - 12.05608232776096 * x * y**4 * w * v - - 10.33378485236654 * x * y**2 * w * v - + 1.033378485236654 * x * w * v - ], - [ - 12.05608232776096 * x**4 * z * w * v - - 10.33378485236654 * x**2 * z * w * v - + 1.033378485236654 * z * w * v - ], - [ - 12.05608232776096 * y**4 * z * w * v - - 10.33378485236654 * y**2 * z * w * v - + 1.033378485236654 * z * w * v - ], - [ - 12.05608232776096 * x * z**4 * w * v - - 10.33378485236654 * x * z**2 * w * v - + 1.033378485236654 * x * w * v - ], - [ - 12.05608232776096 * y * z**4 * w * v - - 10.33378485236654 * y * z**2 * w * v - + 1.033378485236654 * y * w * v - ], - [ - 12.05608232776096 * x * y * w**4 * v - - 10.33378485236654 * x * y * w**2 * v - + 1.033378485236654 * x * y * v - ], - [ - 12.05608232776096 * x * z * w**4 * v - - 10.33378485236654 * x * z * w**2 * v - + 1.033378485236654 * x * z * v - ], - [ - 12.05608232776096 * y * z * w**4 * v - - 10.33378485236654 * y * z * w**2 * v - + 1.033378485236654 * y * z * v - ], - [ - 12.05608232776096 * x * y * z * v**4 - - 10.33378485236654 * x * y * z * v**2 - + 1.033378485236654 * x * y * z - ], - [ - 12.05608232776096 * x * y * w * v**4 - - 10.33378485236654 * x * y * w * v**2 - + 1.033378485236654 * x * y * w - ], - [ - 12.05608232776096 * x * z * w * v**4 - - 10.33378485236654 * x * z * w * v**2 - + 1.033378485236654 * x * z * w - ], - [ - 12.05608232776096 * y * z * w * v**4 - - 10.33378485236654 * y * z * w * v**2 - + 1.033378485236654 * y * z * w - ], - [ - 20.88174713191521 * x**4 * y * z * w * v - - 17.89864039878447 * x**2 * y * z * w * v - + 1.789864039878446 * y * z * w * v - ], - [ - 20.88174713191521 * x * y**4 * z * w * v - - 17.89864039878447 * x * y**2 * z * w * v - + 1.789864039878446 * x * z * w * v - ], - [ - 20.88174713191521 * x * y * z**4 * w * v - - 17.89864039878447 * x * y * z**2 * w * v - + 1.789864039878446 * x * y * w * v - ], - [ - 20.88174713191521 * x * y * z * w**4 * v - - 17.89864039878447 * x * y * z * w**2 * v - + 1.789864039878446 * x * y * z * v - ], - [ - 20.88174713191521 * x * y * z * w * v**4 - - 17.89864039878447 * x * y * z * w * v**2 - + 1.789864039878446 * x * y * z * w - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be <5".format( - order - ) - ) - - elif modal and basis_type == "gkhybrid": - if order == 1: - functionVector = Matrix( - [ - [0.1767766952966368], - [0.3061862178478971 * x], - [0.3061862178478971 * y], - [0.3061862178478971 * z], - [0.3061862178478971 * w], - [0.3061862178478971 * v], - [0.5303300858899105 * x * y], - [0.5303300858899105 * x * z], - [0.5303300858899105 * y * z], - [0.5303300858899105 * w * x], - [0.5303300858899105 * w * y], - [0.5303300858899105 * w * z], - [0.5303300858899105 * v * x], - [0.5303300858899105 * v * y], - [0.5303300858899105 * v * z], - [0.5303300858899105 * v * w], - [0.9185586535436913 * x * y * z], - [0.9185586535436913 * w * x * y], - [0.9185586535436913 * w * x * z], - [0.9185586535436913 * w * y * z], - [0.9185586535436913 * v * x * y], - [0.9185586535436913 * v * x * z], - [0.9185586535436913 * v * y * z], - [0.9185586535436913 * v * w * x], - [0.9185586535436913 * v * w * y], - [0.9185586535436913 * v * w * z], - [1.590990257669731 * w * x * y * z], - [1.590990257669731 * v * x * y * z], - [1.590990257669731 * v * w * x * y], - [1.590990257669731 * v * w * x * z], - [1.590990257669731 * v * w * y * z], - [2.755675960631073 * v * w * x * y * z], - [0.592927061281571 * (w**2 - 0.3333333333333333)], - [1.026979795322186 * (w**2 * x - 0.3333333333333333 * x)], - [1.026979795322186 * (w**2 * y - 0.3333333333333333 * y)], - [1.026979795322186 * (w**2 * z - 0.3333333333333333 * z)], - [1.026979795322186 * (v * w**2 - 0.3333333333333333 * v)], - [1.778781183844713 * (w**2 * x * y - 0.3333333333333333 * x * y)], - [1.778781183844713 * (w**2 * x * z - 0.3333333333333333 * x * z)], - [1.778781183844713 * (w**2 * y * z - 0.3333333333333333 * y * z)], - [1.778781183844713 * (v * w**2 * x - 0.3333333333333333 * v * x)], - [1.778781183844713 * (v * w**2 * y - 0.3333333333333333 * v * y)], - [1.778781183844713 * (v * w**2 * z - 0.3333333333333333 * v * z)], - [ - 3.080939385966558 - * (w**2 * x * y * z - 0.3333333333333333 * x * y * z) - ], - [ - 3.080939385966558 - * (v * w**2 * x * y - 0.3333333333333333 * v * x * y) - ], - [ - 3.080939385966558 - * (v * w**2 * x * z - 0.3333333333333333 * v * x * z) - ], - [ - 3.080939385966558 - * (v * w**2 * y * z - 0.3333333333333333 * v * y * z) - ], - [ - 5.336343551534138 - * (v * w**2 * x * y * z - 0.3333333333333333 * v * x * y * z) - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0] - * interpListND[3].shape[0] - * interpListND[4].shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpListND[4].shape[0]): - for j in range(0, interpListND[3].shape[0]): - for k in range(0, interpListND[2].shape[0]): - for l in range(0, interpListND[1].shape[0]): - for m in range(0, interpListND[0].shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpListND[0].shape[0] - + k * interpListND[0].shape[0] * interpListND[1].shape[0] - + j - * interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0] - + i - * interpListND[0].shape[0] - * interpListND[1].shape[0] - * interpListND[2].shape[0] - * interpListND[3].shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpListND[0][m]) - .subs(y, interpListND[1][l]) - .subs(z, interpListND[2][k]) - .subs(w, interpListND[3][j]) - .subs(v, interpListND[4][i]) - ) - - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be =1".format( - order - ) - ) - - elif modal == False and basis_type == "serendipity": - if order == 1: - functionVector = Matrix( - [ - [ - (v * w) / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - - v / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - - y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - y / 32.0 - - w / 32.0 - - x / 32.0 - - v / 32.0 - - z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - + y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - z / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - - v / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - - y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - y / 32.0 - - w / 32.0 - - x / 32.0 - - v / 32.0 - + z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - x / 32.0 - - w / 32.0 - - v / 32.0 - + y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - - x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - w / 32.0 - - v / 32.0 - + x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - - y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - + y / 32.0 - - z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - - y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - - x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - - w / 32.0 - + x / 32.0 - + y / 32.0 - + z / 32.0 - - (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - - y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - - y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - - (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - + y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - + y / 32.0 - - z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - - (v * z) / 32.0 - + (w * y) / 32.0 - - (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - - (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - - y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - - (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - - (x * z) / 32.0 - - (y * z) / 32.0 - - (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - - (v * x * z) / 32.0 - + (w * x * y) / 32.0 - - (v * y * z) / 32.0 - - (w * x * z) / 32.0 - - (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - - y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - - (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - - (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - + (x * z) / 32.0 - - (y * z) / 32.0 - + (v * w * x) / 32.0 - - (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - + (v * x * z) / 32.0 - - (w * x * y) / 32.0 - - (v * y * z) / 32.0 - + (w * x * z) / 32.0 - - (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - - (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - - x / 32.0 - + y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - - (v * x) / 32.0 - + (v * y) / 32.0 - - (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - - (x * y) / 32.0 - - (x * z) / 32.0 - + (y * z) / 32.0 - - (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - - (v * x * y) / 32.0 - - (v * x * z) / 32.0 - - (w * x * y) / 32.0 - + (v * y * z) / 32.0 - - (w * x * z) / 32.0 - + (w * y * z) / 32.0 - - (x * y * z) / 32.0 - - (v * w * x * y) / 32.0 - - (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - - (v * x * y * z) / 32.0 - - (w * x * y * z) / 32.0 - - (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - [ - v / 32.0 - + w / 32.0 - + x / 32.0 - + y / 32.0 - + z / 32.0 - + (v * w) / 32.0 - + (v * x) / 32.0 - + (v * y) / 32.0 - + (w * x) / 32.0 - + (v * z) / 32.0 - + (w * y) / 32.0 - + (w * z) / 32.0 - + (x * y) / 32.0 - + (x * z) / 32.0 - + (y * z) / 32.0 - + (v * w * x) / 32.0 - + (v * w * y) / 32.0 - + (v * w * z) / 32.0 - + (v * x * y) / 32.0 - + (v * x * z) / 32.0 - + (w * x * y) / 32.0 - + (v * y * z) / 32.0 - + (w * x * z) / 32.0 - + (w * y * z) / 32.0 - + (x * y * z) / 32.0 - + (v * w * x * y) / 32.0 - + (v * w * x * z) / 32.0 - + (v * w * y * z) / 32.0 - + (v * x * y * z) / 32.0 - + (w * x * y * z) / 32.0 - + (v * w * x * y * z) / 32.0 - + 1.0 / 32.0 - ], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, functionVector.shape[0]): - interpMatrix[ - m - + l * interpList.shape[0] - + k * interpList.shape[0] * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - n, - ] = ( - functionVector[n] - .subs(x, interpList[m]) - .subs(y, interpList[l]) - .subs(z, interpList[k]) - .subs(w, interpList[j]) - .subs(v, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be 1 for nodal Serendipity in 5D".format( - order - ) - ) - - else: - raise NameError( - "interpMatrix: Basis {} is not supported!\nSupported basis are currently 'nodal Serendipity', 'modal Serendipity', and 'modal maximal order'".format( - basis_type - ) - ) - - elif dim == 6: - x = Symbol("x") - y = Symbol("y") - z = Symbol("z") - w = Symbol("w") - v = Symbol("v") - u = Symbol("u") - if modal and basis_type == "serendipity": - if order == 0: - functionVector = Matrix([[0.125]]) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, interpList.shape[0]): - for o in range(0, functionVector.shape[0]): - interpMatrix[ - n - + m * interpList.shape[0] - + l * interpList.shape[0] * interpList.shape[0] - + k - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - o, - ] = ( - functionVector[o] - .subs(x, interpList[n]) - .subs(y, interpList[m]) - .subs(z, interpList[l]) - .subs(w, interpList[k]) - .subs(v, interpList[j]) - .subs(u, interpList[i]) - ) - elif order == 1: - functionVector = Matrix( - [ - [0.125], - [0.2165063509461096 * x], - [0.2165063509461096 * y], - [0.2165063509461096 * z], - [0.2165063509461096 * w], - [0.2165063509461096 * v], - [0.2165063509461096 * u], - [0.375 * x * y], - [0.375 * x * z], - [0.375 * y * z], - [0.375 * x * w], - [0.375 * y * w], - [0.375 * z * w], - [0.375 * x * v], - [0.375 * y * v], - [0.375 * z * v], - [0.375 * w * v], - [0.375 * x * u], - [0.375 * y * u], - [0.375 * z * u], - [0.375 * w * u], - [0.375 * v * u], - [0.6495190528383289 * x * y * z], - [0.6495190528383289 * x * y * w], - [0.6495190528383289 * x * z * w], - [0.6495190528383289 * y * z * w], - [0.6495190528383289 * x * y * v], - [0.6495190528383289 * x * z * v], - [0.6495190528383289 * y * z * v], - [0.6495190528383289 * x * w * v], - [0.6495190528383289 * y * w * v], - [0.6495190528383289 * z * w * v], - [0.6495190528383289 * x * y * u], - [0.6495190528383289 * x * z * u], - [0.6495190528383289 * y * z * u], - [0.6495190528383289 * x * w * u], - [0.6495190528383289 * y * w * u], - [0.6495190528383289 * z * w * u], - [0.6495190528383289 * x * v * u], - [0.6495190528383289 * y * v * u], - [0.6495190528383289 * z * v * u], - [0.6495190528383289 * w * v * u], - [1.125 * x * y * z * w], - [1.125 * x * y * z * v], - [1.125 * x * y * w * v], - [1.125 * x * z * w * v], - [1.125 * y * z * w * v], - [1.125 * x * y * z * u], - [1.125 * x * y * w * u], - [1.125 * x * z * w * u], - [1.125 * y * z * w * u], - [1.125 * x * y * v * u], - [1.125 * x * z * v * u], - [1.125 * y * z * v * u], - [1.125 * x * w * v * u], - [1.125 * y * w * v * u], - [1.125 * z * w * v * u], - [1.948557158514986 * x * y * z * w * v], - [1.948557158514986 * x * y * z * w * u], - [1.948557158514986 * x * y * z * v * u], - [1.948557158514986 * x * y * w * v * u], - [1.948557158514986 * x * z * w * v * u], - [1.948557158514986 * y * z * w * v * u], - [3.375 * x * y * z * w * v * u], - ] - ) - interpMatrix = numpy.zeros( - ( - interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - functionVector.shape[0], - ) - ) - for i in range(0, interpList.shape[0]): - for j in range(0, interpList.shape[0]): - for k in range(0, interpList.shape[0]): - for l in range(0, interpList.shape[0]): - for m in range(0, interpList.shape[0]): - for n in range(0, interpList.shape[0]): - for o in range(0, functionVector.shape[0]): - interpMatrix[ - n - + m * interpList.shape[0] - + l * interpList.shape[0] * interpList.shape[0] - + k - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + j - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - + i - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0] - * interpList.shape[0], - o, - ] = ( - functionVector[o] - .subs(x, interpList[n]) - .subs(y, interpList[m]) - .subs(z, interpList[l]) - .subs(w, interpList[k]) - .subs(v, interpList[j]) - .subs(u, interpList[i]) - ) - else: - raise NameError( - "interpMatrix: Order {} is not supported!\nPolynomial order must be 1 for modal Serendipity in 6D".format( - order - ) - ) - - else: - raise NameError( - "interpMatrix: Basis {} is not supported!\nSupported basis are currently 'modal Serendipity' in 6D".format( - basis_type - ) - ) - - else: - raise NameError("interpMatrix: Dimension {} is not supported.".format(dim)) - - return interpMatrix - - -if __name__ == "__main__": - import tables - # set command line options - parser = OptionParser() - parser.add_option( - "-d", "--dimension", action="store", dest="dim", help="specified dimension" - ) - parser.add_option( - "-o", "--order", action="store", dest="order", help="specified polynomial order" - ) - parser.add_option( - "-b", "--basis", action="store", dest="basis", help="specified basis set" - ) - parser.add_option( - "-i", - "--interp", - action="store", - dest="interp", - help="specified number of interpolation points", - ) - parser.add_option( - "-m", - "--modal", - action="store", - dest="modal", - help="set to True for modal basis set", - ) - - (options, args) = parser.parse_args() - - dim = int(options.dim) - order = int(options.order) - basis_type = options.basis - modal = options.modal - interp = int(options.interp) - - interpMatrix = createInterpMatrix(dim, order, basis_type, interp, modal) - fh = tables.open_file("interpMatrix.h5", mode="w") - fh.create_array("/", "interpolation_matrix", interpMatrix) - fh.close() diff --git a/src/postgkyl/data/dg.py b/src/postgkyl/data/dg.py deleted file mode 100644 index 5c569ad6..00000000 --- a/src/postgkyl/data/dg.py +++ /dev/null @@ -1,658 +0,0 @@ -import numpy as np -import os.path -import tables - -from postgkyl.data.computeDerivativeMatrices import createDerivativeMatrix -from postgkyl.data.computeInterpolationMatrices import createInterpMatrix - -# from postgkyl.data.recovData import recovC0Fn, recovC1Fn, recovEdFn - -path = os.path.dirname(os.path.realpath(__file__)) - -num_nodesSerendipity = np.array([ - [1, 2, 3, 4, 5], - [1, 4, 8, 12, 17], - [1, 8, 20, 32, 50], - [1, 16, 48, 80, 136], - [1, 32, 112, 192, 352], - [1, 64, 256, 448, 880]]) - -num_nodesMaximal = np.array([ - [2, 3, 4, 5], - [3, 6, 10, 15], - [4, 10, 20, 35], - [5, 15, 35, 70], - [6, 21, 56, 126], - [7, 28, 84, 210]]) - -num_nodesTensor = np.array([ - [2, 3, 4, 5], - [4, 9, 16, 25], - [8, 27, 64, 125], - [16, 81, 256, 625], - [32, 343, 1024, 3125], - [64, 729, 4096, 15625]]) - -num_nodesGkHybrid = np.array([1, 6, 12, 24, 48]) -num_nodesGkHybridVel = np.array([3, 6]) -num_nodeshybrid = np.array([1, 6, 12, 24, 48]) - - -def _get_basis_p(num_dim, num_comp): - basis, poly_order = None, None - idx = np.argwhere(num_nodesSerendipity[num_dim - 1, :] == num_comp).squeeze() - if idx: - basis = "serendipity" - poly_order = idx - # end - idx = np.argwhere(num_nodesTensor[num_dim - 1, :] == num_comp).squeeze() - if idx: - basis = "tensor" - poly_order = idx + 1 - # end - return basis, poly_order - -def _getnum_nodes(dim, poly_order, basis_type): - if basis_type.lower() == "serendipity": - num_nodes = num_nodesSerendipity[dim - 1, poly_order] - elif basis_type.lower() == "maximal-order": - num_nodes = num_nodesMaximal[dim - 1, poly_order - 1] - elif basis_type.lower() == "tensor": - num_nodes = num_nodesTensor[dim - 1, poly_order - 1] - elif basis_type.lower() == "gkhybrid": - num_nodes = num_nodesGkHybrid[dim - 1] - elif basis_type.lower() == "gkhybrid_vel": - num_nodes = num_nodesGkHybridVel[dim - 1] - elif basis_type.lower() == "hybrid": - num_nodes = num_nodeshybrid[dim - 1] - else: - raise NameError( - "GInterp: Basis '{:s}' is not supported!\n" - "Supported basis are currently 'ns' (Nodal Serendipity)," - " 'ms' (Modal Serendipity), 'mt' (Modal Tensor product)," - " 'mo' (Modal maximal Order), 'gkhybrid' (Modal GkHybrid)," - " 'gkhybrid_vel' (Modal GkHybridVel), and 'hybrid' (Modal hybrid)".format(basis_type) - ) - # end - return num_nodes - -def get_num_basis(dim, poly_order, basis_type) -> int: - # Return the number of nodes for a dimensionality, basis type and poly order. - return _getnum_nodes(dim, poly_order, basis_type) - -def _loadInterpMatrix(dim, poly_order, basis_type, interp, read, modal, c2p=False): - if (interp is not None and read is None) or c2p: - if interp is None: - interp = poly_order + 1 - # end - mat = createInterpMatrix(dim, poly_order, basis_type, interp, modal, c2p) - return mat - elif basis_type == "tensor": - mat = createInterpMatrix(dim, poly_order, "tensor", poly_order + 1, True, c2p) - return mat - elif basis_type == "gkhybrid": - mat = createInterpMatrix(dim, poly_order, "gkhybrid", poly_order + 1, True, c2p) - return mat - elif basis_type == "gkhybrid_vel": - mat = createInterpMatrix(dim, poly_order, "gkhybrid_vel", poly_order + 1, True, c2p) - return mat - elif basis_type == "hybrid": - mat = createInterpMatrix(dim, poly_order, "hybrid", poly_order + 1, True, c2p) - return mat - else: - # Load interpolation matrix from the pre-computed HDF5 file. - varid = "xformMatrix%i%i" % (dim, poly_order) - if modal == False and basis_type.lower() == "serendipity": - fileName = path + "/xformMatricesNodalSerendipity.h5" - elif modal and basis_type.lower() == "serendipity": - fileName = path + "/xformMatricesModalSerendipity.h5" - - elif modal and basis_type.lower() == "maximal-order": - fileName = path + "/xformMatricesModalMaximal.h5" - else: - raise NameError( - "GInterp: Basis {:s} is not supported!\n" - "Supported basis are currently 'ns' (Nodal Serendipity), " - "'ms' (Modal Serendipity), and 'mo' (Modal Maximal Order)".format(basis_type) - ) - # end - fh = tables.open_file(fileName) - mat = fh.root.matrices._v_children[varid].read() - fh.close() - return mat.transpose() - # end - - -def _loadDerivativeMatrix(dim, poly_order, basis_type, interp, read, modal=True): - if interp is not None and read is None: - mat = createDerivativeMatrix(dim, poly_order, basis_type, interp, modal) - return mat - else: - interp = poly_order + 1 - mat = createDerivativeMatrix(dim, poly_order, basis_type, interp, modal) - return mat - # end - - -def _makeMesh(num_interp, Xc, xlo=None, xup=None, gridType=None): - nx = Xc.shape[0] - 1 # expecting nodal mesh - meshOut = np.zeros(num_interp * nx + 1) - if gridType is None or gridType == "uniform": - if xlo is None or xup is None: - xlo = Xc[0] - xup = Xc[-1] - # end - meshOut = np.linspace(xlo, xup, num_interp*nx + 1) - elif gridType == "mapped": - # subdivide every cell in Xc into num_interp cells. - for i in range(nx): - dx = (Xc[i + 1] - Xc[i]) / num_interp - for j in range(num_interp): - meshOut[i*num_interp + j] = Xc[i] + j*dx - # end - # end - # add the last node. - dx = (Xc[-1] - Xc[-2]) / num_interp - meshOut[nx*num_interp] = Xc[nx - 1] + num_interp*dx - # end - return meshOut - - -def _make1Dgrids(num_interp, Xc, num_dims, gridType=None): - # build a list of 1D arrays, each containing the grid in that dimension. - gridOut = list() - if gridType is None or gridType == "uniform": - gridOut = [_makeMesh(num_interp[d], Xc[d]) for d in range(num_dims)] - elif gridType == "mapped": - # back out 1D arrays from Xc. - for d in range(num_dims): - currSlices = [0] * num_dims - currSlices[-1 - d] = np.s_[:] - gridOut.append(_makeMesh(num_interp[d], Xc[d][tuple(currSlices)], gridType=gridType)) - # end - # end - return gridOut - - -def _interpOnMesh(cMat, qIn, nInterpIn, basis_type, c2p=False): - numCells = np.array(qIn.shape) - # last entry is indexing nodes, get rid of it - numCells = numCells[:-1] - num_dims = int(len(numCells)) - num_interp = np.array([max(nInterpIn, 2)] * num_dims) - if basis_type == "gkhybrid": - # 1x1v, 1x2v, 2x2v, 3x2v cases, with p=2 in the first velocity dim. - vpardir = (1 if (num_dims == 2 or num_dims == 3) else - (2 if num_dims == 4 else - (3 if num_dims == 5 else 99 ) ) ) - num_interp[vpardir] = nInterpIn + 1 - # end - if basis_type == "gkhybrid_vel": - # 1v, 2v with p=2 in the first velocity dim. - vpardir = 0 - num_interp[vpardir] = nInterpIn + 1 - # end - if basis_type == "hybrid": - num_interp[-1] = nInterpIn + 1 - # end - if c2p: - qOut = np.zeros(numCells*(num_interp - 1) + 1, np.float64) - else: - qOut = np.zeros(numCells*num_interp, np.float64) - # end - # move the node index from last to the first - qIn = np.moveaxis(qIn, -1, 0) - # Main loop - for n in range(np.prod(num_interp)): - # https://docs.scipy.org/doc/numpy/reference/generated/numpy.tensordot.html - temp = np.tensordot(cMat[n, :], qIn, axes=1) - # decompose n to i,j,k,... indices based on the number of dimensions - startIdx = np.unravel_index(n, num_interp, order="F") - # define multi-D qOut slices - if c2p: - idxs = [slice(int(startIdx[i]), int(numCells[i]*(num_interp[i] - 1) + startIdx[i]), - num_interp[i] - 1) - for i in range(num_dims)] - else: - idxs = [slice(int(startIdx[i]), int(numCells[i]*num_interp[i]), num_interp[i]) - for i in range(num_dims)] - # end - qOut[tuple(idxs)] = temp - # end - return np.array(qOut) - - -class GInterp(object): - """Postgkyl base class for DG data manipulation. - - This class should not be used on its own! Currently supported - child classes are: - - GInterpNodal - - GInterpModal - - Init Args: - data (GData): Data to work with - num_nodes (int): Number of nodes - """ - - def __init__(self, data, num_nodes): - self.data = data - self.num_nodes = num_nodes - self.numEqns = data.get_num_comps() / num_nodes - self.num_dims = data.get_num_dims() - self.Xc = data.get_grid() - self.gridType = data.get_grid_type() - - def _getRawNodal(self, component): - q = self.data.get_values() - numEqns = self.numEqns - shp = [q.shape[i] for i in range(self.num_dims)] - shp.append(self.num_nodes) - rawData = np.zeros(shp, np.float64) - for n in range(self.num_nodes): - rawData[..., n] = q[..., int(component + n * numEqns)] - # end - return rawData - - def _getRawModal(self, component): - q = self.data.get_values() - shp = [q.shape[i] for i in range(self.num_dims)] - shp.append(self.num_nodes) - rawData = np.zeros(shp, np.float64) - lo = int(component * self.num_nodes) - up = int(lo + self.num_nodes) - rawData = q[..., lo:up] - return rawData - - -class GInterpNodal(GInterp): - """Postgkyl class for nodal DG data manipulation. - - After the initializations, GInterpNodal object provides the - interpolate and differentiate methods. These returns grid and - values by default but could be used to directly push to the GData - stack with the stack=True flag. - - Parent: GInterp - - Init Args: - data (GData): Data to work with - poly_order (int): Order of the polynomial approximation - basis (str): Specify the basis. Currently supported is the - nodal Serendipity 'ns' - num_interp (int): Specify number of points on which to - interpolate (default: poly_order + 1) - read - - Example: - import postgkyl - data = postgkyl.GData('file.h5') - dg = postgkyl.GInterpNodal(data, 2, 'ns') - grid, values = dg.interpolate() - """ - - def __init__(self, data, poly_order, basis_type, num_interp=None, read=None): - self.num_dims = data.get_num_dims() - self.poly_order = poly_order - self.basis_type = basis_type - if basis_type == "ns": - self.basis_type = "serendipity" - # end - - self.num_interp = num_interp - self.read = read - num_nodes = _getnum_nodes(self.num_dims, self.poly_order, self.basis_type) - GInterp.__init__(self, data, num_nodes) - - def interpolate(self, comp=0, overwrite=False, stack=False): - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - cMat = _loadInterpMatrix(self.num_dims, self.poly_order, self.basis_type, - self.num_interp, self.read, False) - if isinstance(comp, int): - q = self._getRawNodal(comp) - values = _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis] - elif isinstance(comp, tuple): - q = self._getRawNodal(comp[0]) - values = _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis] - for c in comp[1:]: - q = self._getRawNodal(c) - values = np.append(values, - _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis], - axis=-1) - # end - elif isinstance(comp, slice): - q = self._getRawNodal(comp.start) - values = _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis] - for c in range(comp.start + 1, comp.stop): - q = self._getRawNodal(c) - values = np.append( - values, - _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis], - axis=-1, - ) - # end - # end - - num_interp = [int(round(cMat.shape[0] ** (1.0 / self.num_dims)))] * self.num_dims - grid = _make1Dgrids(num_interp, self.Xc, self.num_dims) - if overwrite: - self.data.push(grid, values) - else: - return grid, values - # end - - def differentiate(self, direction, comp=0, overwrite=False, stack=False): - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - q = self._getRawNodal(comp) - cMat = _loadDerivativeMatrix(self.num_dims, self.poly_order, self.basis_type, - self.num_interp, self.read, False) - if direction is not None: - values = ( - _interpOnMesh(cMat[:, :, direction], q, self.num_interp, self.basis_type) - * 2 - / (self.Xc[direction][1] - self.Xc[direction][0])) - values = values[..., np.newaxis] - else: - values = np.zeros(q.shape, self.num_dims) - for i in range(self.num_dims): - values[:, i] = _interpOnMesh(cMat[:, :, i], q, self.num_interp, self.basis_type) - values[:, i] *= 2 / (self.Xc[i][1] - self.Xc[i][0]) - # end - # end - - num_interp = [int(round(cMat.shape[0] ** (1.0 / self.num_dims)))] * self.num_dims - grid = _make1Dgrids(num_interp, self.Xc, self.num_dims) - if overwrite: - self.data.push(grid, values) - else: - return grid, values - # end - - -class GInterpModal(GInterp): - """Postgkyl class for modal DG data manipulation. - - After the initializations, GInterpModal object provides the - interpolate and differentiate methods. These returns grid and - values by default but could be used to directly push to the GData - stack with the stack=True flag. - - Parent: GInterp - - Init Args: - data (GData): Data to work with - poly_order (int): Order of the polynomial approximation - basis (str): Specify the basis. Currently supported are the - modal Serendipity 'ms' and the maximal order basis 'mo' - num_interp (int): Specify number of points on which to - interpolate (default: poly_order + 1) - read - - Example: - import postgkyl - data = postgkyl.GData('file.bp') - dg = postgkyl.GInterpModal(data, 2, 'ms') - grid, values = dg.interpolate() - """ - - def __init__(self, data, poly_order=None, basis_type=None, num_interp=None, - periodic=False, read=None): - self.num_dims = data.get_num_dims() - if poly_order is not None: - self.poly_order = poly_order - elif data.ctx.get("poly_order"): - self.poly_order = data.ctx["poly_order"] - else: - raise ValueError( - "GInterpNodal: polynomial order is neither specified nor stored in the output file") - # end - if basis_type: - if basis_type == "ms": - self.basis_type = "serendipity" - elif basis_type == "mo": - self.basis_type = "maximal-order" - elif basis_type == "mt": - self.basis_type = "tensor" - elif basis_type == "gkhyb": - self.basis_type = "gkhybrid" - elif basis_type == "gkhyb_vel": - self.basis_type = "gkhybrid_vel" - elif basis_type == "pkpmhyb": - self.basis_type = "hybrid" - # end - elif data.ctx.get("basis_type"): - self.basis_type = data.ctx["basis_type"] - else: - raise ValueError( - "GInterpModal: basis type is neither specified nor stored in the output file") - # end - - # PKPM hybrid base expects 2+ dimensions with the last one being - # the parallel velocity. This allows to specify 'pkpmhyb' basis - # and work with 1x1v and 1x data simulataneously. - if self.num_dims == 1 and self.basis_type == "hybrid": - self.basis_type = "serendipity" - # end - - self.periodic = periodic - - # XXX This was introduced with the c2p but I can't see the importance of the extra - # condition and seem to unecessarily limit the capabilities. The c2p test cases - # still seems to produce correct results. -- P.C. - # if num_interp is not None and self.poly_order > 1: - if num_interp: - self.num_interp = num_interp - else: - self.num_interp = self.poly_order + 1 - # end - self.read = read - num_nodes = _getnum_nodes(self.num_dims, self.poly_order, self.basis_type) - GInterp.__init__(self, data, num_nodes) - - def interpolate(self, comp=0, overwrite=False, stack=False): - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - cMat = _loadInterpMatrix(self.num_dims, self.poly_order, self.basis_type, - self.num_interp, self.read, True) - if isinstance(comp, int): - q = self._getRawModal(comp) - values = _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis] - elif isinstance(comp, tuple): - q = self._getRawModal(comp[0]) - values = _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis] - for c in comp[1:]: - q = self._getRawModal(c) - values = np.append(values, - _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis], - axis=-1) - # end - elif isinstance(comp, slice): - q = self._getRawModal(comp.start) - values = _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis] - for c in range(comp.start + 1, comp.stop): - q = self._getRawModal(c) - values = np.append(values, - _interpOnMesh(cMat, q, self.num_interp, self.basis_type)[..., np.newaxis], - axis=-1) - # end - # end - if self.data.ctx["grid_type"] == "c2p": - q = self.data.get_grid() - num_comp = q[0].shape[-1] - basis, poly_order = _get_basis_p(self.num_dims, num_comp) - cMat = _loadInterpMatrix(self.num_dims, poly_order, basis, self.num_interp, - self.read, True, True) - grid = [] - for d in range(self.num_dims): - grid.append(_interpOnMesh(cMat, q[d], self.num_interp + 1, basis, True)) - # end - else: - if self.basis_type == "gkhybrid": - # 1x1v, 1x2v, 2x2v, 3x2v cases, with p=2 in the first velocity dim. - vpardir = (1 if (self.num_dims == 2 or self.num_dims == 3) - else (2 if self.num_dims == 4 else (3 if self.num_dims == 5 else 99))) - num_interp = [self.num_interp] * self.num_dims - num_interp[vpardir] = self.num_interp + 1 - elif self.basis_type == "gkhybrid_vel": - # 1v, 2v, with p=2 in the first velocity dim. - vpardir = 0 - num_interp = [self.num_interp] * self.num_dims - num_interp[vpardir] = self.num_interp + 1 - elif self.basis_type == "hybrid": - num_interp = [self.num_interp] * self.num_dims - num_interp[-1] = self.num_interp + 1 - else: - num_interp = [int(round(cMat.shape[0] ** (1.0 / self.num_dims)))] * self.num_dims - # end - - grid = _make1Dgrids(num_interp, self.Xc, self.num_dims, None) - if self.data.ctx["grid_type"] == "c2p_vel": - num_cdim = self.data.ctx["num_cdim"] - num_vdim = self.data.ctx["num_vdim"] - q = self.data.get_grid() - num_comp = q[-1].shape[-1] - basis, poly_order = _get_basis_p(1, num_comp) - for d in range(num_vdim): - cMat = _loadInterpMatrix(1, poly_order, basis, num_interp[num_cdim + d], - self.read, True, True) - grid[num_cdim + d] = _interpOnMesh(cMat, q[num_cdim + d], - num_interp[num_cdim + d] + 1, basis, True) - # end - # end - # end - - if overwrite: - self.data.push(grid, values) - else: - return grid, values - # end - - def interpolateGrid(self, overwrite=False): - if self.data.ctx["grid_type"] == "c2p": - q = self.data.get_grid() - num_comp = q[0].shape[-1] - basis, poly_order = _get_basis_p(self.num_dims, num_comp) - cMat = _loadInterpMatrix(self.num_dims, poly_order, basis, self.num_interp, - self.read, True, True) - grid = [] - for d in range(self.num_dims): - grid.append(_interpOnMesh(cMat, q[d], self.num_interp, self.basis_type, True)) - # end - elif self.data.ctx["grid_type"] == "c2p_vel": - q = self.data.get_grid() - else: - num_interp = [self.num_interp] * self.num_dims - grid = _make1Dgrids(num_interp, self.Xc, self.num_dims, self.gridType) - # end - - if overwrite: - self.data.set_grid(grid) - else: - return grid - # end - - def differentiate(self, direction=None, comp=0, overwrite=False, stack=False): - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - q = self._getRawModal(comp) - cMat = _loadDerivativeMatrix(self.num_dims, self.poly_order, self.basis_type, - self.num_interp, self.read, True) - if direction is not None: - values = (_interpOnMesh(cMat[:, :, direction], q, self.num_interp, self.basis_type)*2 - / (self.Xc[direction][1] - self.Xc[direction][0])) - values = values[..., np.newaxis] - else: - values = _interpOnMesh(cMat[..., 0], q, self.num_interp, self.basis_type) - values /= self.Xc[0][1] - self.Xc[0][0] - values = values[..., np.newaxis] - for i in range(1, self.num_dims): - values = np.append(values, - _interpOnMesh(cMat[..., i], q, self.num_interp, self.basis_type)[..., np.newaxis], - axis=self.num_dims) - values[..., i] *= 2 / (self.Xc[i][1] - self.Xc[i][0]) - # end - # end - - num_interp = [int(round(cMat.shape[0] ** (1.0 / self.num_dims)))] * self.num_dims - grid = _make1Dgrids(num_interp, self.Xc, self.num_dims, self.gridType) - if overwrite: - self.data.push(grid, values) - else: - return grid, values - # end - - # def recovery(self, comp=0, c1=False, overwrite=False, stack=False): - # if stack: - # overwrite = stack - # print( - # "Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'" - # ) - # # end - # if isinstance(comp, int): - # q = self._getRawModal(comp) - # else: - # raise ValueError("recovery: only 'int' comp implemented so far") - # # end - # if self.num_dims > 1: - # raise ValueError("recovery: only 1D implemented so far") - # # end - - # if self.num_interp is not None: - # N = self.num_interp - # else: - # N = 100 - # # end - - # numCells = self.data.get_num_cells() - # grid = [ - # np.linspace(self.Xc[int(d)][0], self.Xc[int(d)][-1], int(numCells * N + 1)) - # for d in range(self.num_dims) - # ] - - # values = np.zeros(numCells * N) - # dx = self.Xc[0][1] - self.Xc[0][0] - - # xC = np.linspace(-1, 1, N, endpoint=False) * dx / 2 - # xL = np.linspace(-1, 0, N, endpoint=False) * dx - # xR = np.linspace(0, 1, N, endpoint=False) * dx - - # if self.periodic: - # if c1: - # values[:N] = recovC1Fn[self.poly_order - 1](xC, q[0], q[-1], q[1], dx) - # values[-N:] = recovC1Fn[self.poly_order - 1](xC, q[-1], q[-2], q[0], dx) - # else: - # values[:N] = recovC0Fn[self.poly_order - 1](xC, q[0], q[-1], q[1], dx) - # values[-N:] = recovC0Fn[self.poly_order - 1](xC, q[-1], q[-2], q[0], dx) - # # end - # else: - # values[:N] = recovEdFn[self.poly_order - 1](xL, q[0], q[1], dx) - # values[-N:] = recovEdFn[self.poly_order - 1](xR, q[-2], q[-1], dx) - # # end - # for j in range(1, numCells[0] - 1): - # if c1: - # values[j * N : (j + 1) * N] = recovC1Fn[self.poly_order - 1]( - # xC, q[j], q[j - 1], q[j + 1], dx - # ) - # else: - # values[j * N : (j + 1) * N] = recovC0Fn[self.poly_order - 1]( - # xC, q[j], q[j - 1], q[j + 1], dx - # ) - # # end - # # end - - # values = values[..., np.newaxis] - # if overwrite: - # self.data.push(grid, values) - # else: - # return grid, values - # # end diff --git a/src/postgkyl/data/flash_h5_reader.py b/src/postgkyl/data/flash_h5_reader.py deleted file mode 100644 index c6fe80c4..00000000 --- a/src/postgkyl/data/flash_h5_reader.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Module including FLASH reader class""" - -import math -import numpy as np -import tables -from typing import Tuple - -# FLASH variable names -# dens : the density in g/cc -# tele : the electron temperature in K -# tion : same but for the ions -# velx : the fluid velocity in x direction -# vely : the fluid velocity in y direction -# temp : the overall fluid temperature in K -# pres : the pressure in dyn/cm^2 -# ye -# sumy -# -# The last two variables are used to retrieve the ion and electron -# density in /cc: -# n_ele = ye * Na * dens -# n_ion = sumy * Na * dens -# where Na=6.02e23 is the Avogadro number. -# -# The average ionisation Z' and average atomic mass A' can be found by: -# Z' = ye/sumy -# A' = 1/sumy - - -class FlashH5Reader(object): - """Provides a framework to read FLASH h5 output""" - - def __init__(self, file_name: str, var_name: str, ctx: dict = None, **kwargs) -> None: - self._file_name = file_name - self.var_name = var_name - - self.ctx = ctx - - def is_compatible(self) -> bool: - out = False - try: - fh = tables.open_file(self._file_name, "r") - except: - return False - # end - if "coordinates" in fh.root: - out = True - # end - fh.close() - return out - - def _read_frame(self) -> tuple: - fh = tables.open_file(self._file_name, "r") - coord = fh.root["coordinates"].read().transpose() - bsize = fh.root["block size"].read().transpose() - ntype = fh.root["node type"].read().transpose() - bdata = fh.root[self.var_name].read().transpose() - - nxb, nyb, _, N = bdata.shape - res = bsize.min(axis=1) - lower = (coord - bsize / 2).min(axis=1) - upper = (coord + bsize / 2).max(axis=1) - - nxax = math.floor((upper[0] - lower[0]) / (res[0] / nxb)) - nyax = math.floor((upper[1] - lower[1]) / (res[1] / nyb)) - data = np.zeros((nxax, nyax)) - for b in range(N): - if ntype[b] == 1: - mult = np.ceil(bsize[:, b] / res) - idxx = math.floor((coord[0, b] - bsize[0, b] / 2 - lower[0]) / res[0] * nxb) - idxy = math.floor((coord[1, b] - bsize[1, b] / 2 - lower[1]) / res[1] * nyb) - for i in range(nxb): - for j in range(nyb): - data[ - idxx + i * int(mult[0]) : idxx + (i + 1) * int(mult[0]) + 1, - idxy + j * int(mult[1]) : idxy + (j + 1) * int(mult[1]) + 1, - ] = bdata[i, j, 0, b] - # end - # end - # end - # end - fh.close() - return data.shape, lower[:2], upper[:2], data[..., np.newaxis] - - # ---- Exposed functions ---- - def get_data(self) -> Tuple[np.ndarray, np.ndarray]: - cells, lower, upper, data = self._read_frame() - num_dims = len(cells) - grid = [np.linspace(lower[d], upper[d], cells[d] + 1) for d in range(num_dims)] - - return grid, data diff --git a/src/postgkyl/data/gdata.py b/src/postgkyl/data/gdata.py deleted file mode 100644 index 12d4e1aa..00000000 --- a/src/postgkyl/data/gdata.py +++ /dev/null @@ -1,661 +0,0 @@ -"""Module including Gkeyll data class""" - -from typing import Literal, Tuple -import numpy as np -import shutil - -try: - import adios2 - has_adios = True -except ModuleNotFoundError: - has_adios = False -# end - -from postgkyl.data.gkyl_reader import GkylReader -from postgkyl.data.gkyl_adios_reader import GkylAdiosReader -from postgkyl.data.gkyl_h5_reader import GkylH5Reader -from postgkyl.data.flash_h5_reader import FlashH5Reader -import postgkyl.utils.gkeyll_enums as gkenums - - -class GData(object): - """Provides interface to (not only) Gkeyll output data. - - GData serves as a baseline interface to Gkeyll data. It is used for - loading Gkeyll data and serves is input to many Postgkyl - functions. Represents a dataset in the Postgkyl command line mode. - - Examples: - import postgkyl as pg - data = pg.GData('file.gkyl', comp=1) - - """ - - def __init__(self, file_name: str = "", - comp: int | str | None = None, - z0: int | str | None = None, z1: int | str | None = None, - z2: int | str | None = None, z3: int | str | None = None, - z4: int | str | None = None, z5: int | str | None = None, - var_name: str = "CartGridField", - tag: str = "default", label: str = "", - ctx: dict | None = None, - comp_grid: bool = False, mapc2p_name: str = "", mapc2p_vel_name: str = "", - reader_name: str = "", load: bool = True, click_mode: bool = False): - """Initializes the Data class with a Gkeyll output file. - - Args: - fileName: str - The name of Gkeyll output file. Currently supported are 'h5', - ADIOS 'bp', and binary 'gkyl' files. Can be ommited for empty - class. - comp: int or 'int:int' - Load only the specified component index or a slice of - idices. Supported only for the ADIOS 'bp' files. - z0 - z5: int or 'int:int' - Load only the specified index or a slice of - idices in a direction. Supported only for the ADIOS 'bp' files. - var_name: str - Specify custom ADIOS variable name (default is 'CartGridField'). - tag: str - Specify dataset tag for use in the command line mode. - label: str - Specify dataset label for use in the command line mode. - ctx: dict - Copy content of the specified ctx dictionary. - comp_grid: bool - A flag to ignore grid mapping. - mapc2p_name: str - The name of the file containg the c2p mapping. - mapc2p_vel_name: str - The name of the file containg the c2p mapping just for velocity. - reader_name: str - Reader can be specified to bypass the automatic selection. - load: bool = True - Automatically the data to memory; when set to False, data can be loaded later - using the load() method. - click_mode: bool = False - Enables command-line behavior like prompting when a - var_name is either missing or doesn't match any available. - """ - self._grid = None - self._values = None # (N+1)D narray of values - - # Context dictionary to store metadata, filled by the reader. - self.ctx = {} - - # Allow to copy input context variable - if ctx: - for key in ctx: - self.ctx[key] = ctx[key] - - self._tag = tag - self._comp_grid = comp_grid # flag to disregard the mapped grid - self._label = "" - self._custom_label = label - self._var_name = var_name - self._file_name = str(file_name) - self._mapc2p_name = mapc2p_name - self._mapc2p_vel_name = mapc2p_vel_name - self.color = None - - self._neighbors = [] - - self._status = True - - zs = (z0, z1, z2, z3, z4, z5) - - readers = { - "gkyl": GkylReader, - "adios": GkylAdiosReader, - "h5": GkylH5Reader, - "flash": FlashH5Reader, - } - if self._file_name: - reader_set = False - if reader_name in readers: - # Keep only the user-specified reader - reader = readers[reader_name] - readers.clear() - readers[reader_name] = reader - # end - for key, rd in readers.items(): - self._reader = rd(file_name=self._file_name, ctx=self.ctx, var_name=var_name, - c2p=mapc2p_name, c2p_vel=mapc2p_vel_name, axes=zs, comp=comp, - click_mode=click_mode) - if self._reader.is_compatible(): - reader_set = True - break - # end - # end - if not reader_set: - raise NameError(f"'file_name' was specified ({self._file_name}) but cannot be read with {list(readers)}") - # end - - self._reader.preload() - if load: - self._grid, self._values = self._reader.load() - # end - # end - - # ---- Tag ---- - def get_tag(self) -> str: - return self._tag - - def set_tag(self, tag: str = "") -> None: - if tag: - self._tag = tag - # end - - tag = property(get_tag, set_tag) - - # ---- Label ---- - def get_label(self) -> str: - if self._custom_label: - return self._custom_label - else: - return self._label - # end - - def set_label(self, label: str) -> None: - self._label = label - - label = property(get_label, set_label) - - def get_custom_label(self): - return self._custom_label - - # ---- Status ---- - def activate(self) -> None: - self._status = True - - def deactivate(self) -> None: - self._status = False - - def get_status(self) -> bool: - return self._status - - status = property(get_status) - - # ---- File name ---- - def get_file_name(self) -> str: - return self._file_name - - # ---- Input file ---- - def get_input_file(self) -> str: - if not has_adios: - raise ModuleNotFoundError("ADIOS2 is not installed") - # end - - fh = adios2.open(self._file_name, "rra") - input_file = fh.read_attribute_string("inputfile")[0] - fh.close() - return input_file - - # ---- Number of Cells ---- - def get_num_cells(self) -> np.ndarray: - if self.ctx.get("cells") is not None: - return np.array(self.ctx["cells"]) - elif self._values is not None: - num_dims = len(self._values.shape) - 1 - cells = np.zeros(num_dims, np.int32) - for d in range(num_dims): - cells[d] = int(self._values.shape[d]) - # end - return cells - else: - return 0 - # end - - num_cells = property(get_num_cells) - - # ---- Number of Components ---- - def get_num_comps(self) -> int: - if self.ctx.get("num_comps"): - return self.ctx["num_comps"] - elif self._values is not None: - return int(self._values.shape[-1]) - else: - return 0 - # end - - num_comps = property(get_num_comps) - - # ---- Number of Dimensions ----- - def get_num_dims(self, squeeze: bool = False) -> int: - if self.ctx.get("cells") is not None: - num_dims = len(self.ctx["cells"]) - elif self._values is not None: - num_dims = int(len(self._values.shape) - 1) - else: - return 0 - # end - if squeeze: - cells = self.get_num_cells() - for d in range(num_dims): - if cells[d] == 1: - num_dims = num_dims - 1 - # end - # end - # end - return num_dims - - num_dims = property(get_num_dims) - - # ---- Grid Bounds ---- - def get_bounds(self) -> Tuple[np.ndarray, np.ndarray]: - if "lower" in self.ctx.keys() and "upper" in self.ctx.keys(): - return self.ctx["lower"], self.ctx["upper"] - elif self._grid is not None: - num_dims = len(self._values.shape) - 1 - lo, up = np.zeros(num_dims), np.zeros(num_dims) - for d in range(num_dims): - lo[d] = self._grid[d].min() - up[d] = self._grid[d].max() - # end - return lo, up - else: - return None, None - # end - - bounds = property(get_bounds) - - # ---- Grid and Values ---- - def get_grid(self) -> list: - return self._grid - - def set_grid(self, grid: list) -> None: - self._grid = grid - num_dims = self.get_num_dims() - lo, up = np.zeros(num_dims), np.zeros(num_dims) - for d in range(num_dims): - lo[d] = self._grid[d].min() - up[d] = self._grid[d].max() - self.ctx["lower"] = lo - self.ctx["upper"] = up - - grid = property(get_grid, set_grid) - - def get_grid_type(self) -> str: - return self.ctx["grid_type"] - - def get_values(self) -> np.ndarray: - return self._values - - def set_values(self, values) -> None: - self._values = values - if "cells" not in self.ctx or not np.array_equal(values.shape[:-1], self.ctx["cells"]): - self.ctx["cells"] = values.shape[:-1] - if "num_comps" not in self.ctx or values.shape[-1] != self.ctx["num_comps"]: - self.ctx["num_comps"] = values.shape[-1] - - values = property(get_values, set_values) - - def push(self, grid, values): - self.set_values(values) - self.set_grid(grid) - return self - - # ---- Neighboring Blocks ---- - def set_neighbors(self, dataspace): - data_list = list(dataspace) - num_dims = self.get_num_dims() - for dim in range(num_dims): - self._neighbors.append([None, None]) - for data in data_list: - if num_dims == 1: - if np.isclose(self.get_grid()[dim][0], data.get_grid()[dim][-1]): - self._neighbors[dim][0] = data - elif np.isclose(self.get_grid()[dim][-1], data.get_grid()[dim][0]): - self._neighbors[dim][1] = data - elif num_dims == 2: - if np.isclose(self.get_grid()[dim][0], data.get_grid()[dim][-1]) and np.isclose(self.get_grid()[not dim][0], data.get_grid()[not dim][0]): - self._neighbors[dim][0] = data - elif np.isclose(self.get_grid()[dim][-1], data.get_grid()[dim][0]) and np.isclose(self.get_grid()[not dim][0], data.get_grid()[not dim][0]): - self._neighbors[dim][1] = data - elif num_dims == 3: - rem_dims = list(range(num_dims)).remove(dim) - if np.isclose(self.get_grid()[dim][0], data.get_grid()[dim][-1]) and np.isclose(self.get_grid()[rem_dims[0]][0], data.get_grid()[rem_dims[0]][0]) and np.isclose(self.get_grid()[rem_dims[1]][0], data.get_grid()[rem_dims[1]][0]): - self._neighbors[dim][0] = data - elif np.isclose(self.get_grid()[dim][0], data.get_grid()[dim][-1]) and np.isclose(self.get_grid()[rem_dims[0]][0], data.get_grid()[rem_dims[0]][0]) and np.isclose(self.get_grid()[rem_dims[1]][0], data.get_grid()[rem_dims[1]][0]): - self._neighbors[dim][1] = data - # end - # end - # end - - def _dict_has_key_from_group(self, dict_in, group_members_in): - """ - Check if a dictionary with key-value pairs, where the key is the name of a group and - the value a list of group members (as strings), has a member from a given - group. - """ - return not dict_in.keys().isdisjoint(group_members_in) - - # ---- Info ----- - def info(self) -> str: - """Prints GData object information. - - Prints time (only when available), number of components, dimension - spans, extremes for a GData object. - - Args: - none - - Returns: - output: str - A list of strings with the informations - """ - values = self.values - num_comps = self.num_comps - num_dims = self.num_dims - num_cells = self.num_cells - lower, upper = self.bounds - - # Groups of metadata. - info_groups = { - "time_info" : ["time","frame"], - "grid_info" : ["lower","upper","cells","grid_type"], - "basis_info" : ["poly_order","basis_type","is_modal","num_comps"], - "build_info" : ["changeset","builddate"], - "geometry_info": ["geometry_type", "geqdsk_sign_convention", "is_multib", "topo_file", "half_domain", "geqdsk_file"], - "species_info": ["mass","charge","adiabatic_gamma","vdim"], - } - - output = "" - - printed_keys = [] - - if "time" in self.ctx.keys(): - printed_keys.append("time") - output += f"├─ Time: {self.ctx['time']:e}\n" - # end - - if "frame" in self.ctx.keys(): - printed_keys.append("frame") - output += f"├─ Frame: {self.ctx['frame']:d}\n" - # end - - output += f"├─ Number of components: {num_comps:d}\n" - output += f"├─ Number of dimensions: {num_dims:d}\n" - if self._dict_has_key_from_group(self.ctx, info_groups["grid_info"]): - output += f"├─ Grid: ({self.get_grid_type():s})\n" - if "lower" in self.ctx.keys() and "upper" in self.ctx.keys() and "cells" in self.ctx.keys(): - for d in range(num_dims - 1): - output += f"│ ├─ Dim {d:d}: Num. cells: {num_cells[d]:d}; " - output += f"Lower: {lower[d]:e}; Upper: {upper[d]:e}\n" - # end - # end - - output += f"│ └─ Dim {num_dims - 1:d}: Num. cells: {num_cells[-1]:d}; " - output += f"Lower: {lower[-1]:e}; Upper: {upper[-1]:e}" - # end - - if values is not None: - maximum = np.nanmax(values) - max_idx = np.unravel_index(np.nanargmax(values), values.shape) - minimum = np.nanmin(values) - min_idx = np.unravel_index(np.nanargmin(values), values.shape) - output += f"\n├─ Maximum: {maximum:e} at {str(max_idx[:num_dims]):s}" - if num_comps > 1: - output += f" component {max_idx[-1]:d}\n" - else: - output += "\n" - # end - output += f"├─ Minimum: {minimum:e} at {str(min_idx[:num_dims]):s}" - if num_comps > 1: - output += f" component {min_idx[-1]:d}" - # end - # end - - if self._dict_has_key_from_group(self.ctx, info_groups["basis_info"]): - output += "\n├─ DG info:" - if "poly_order" in self.ctx.keys(): - printed_keys.append("poly_order") - output += f"\n│ ├─ Polynomial Order: {self.ctx['poly_order']:d}" - # end - if "basis_type" in self.ctx.keys(): - printed_keys.append("basis_type") - if self.ctx["is_modal"]: - output += f"\n│ └─ Basis Type: {self.ctx['basis_type']:s} (modal)" - else: - output += f"\n│ └─ Basis Type: {self.ctx['basis_type']:s}" - # end - # end - # end - - if self._dict_has_key_from_group(self.ctx, info_groups["build_info"]): - output += "\n├─ Created with Gkeyll:" - if "changeset" in self.ctx.keys(): - printed_keys.append("changeset") - output += f"\n│ ├─ Changeset: {self.ctx['changeset']:s}" - # end - if "builddate" in self.ctx.keys(): - printed_keys.append("builddate") - output += f"\n│ └─ Build Date: {self.ctx['builddate']:s}" - # end - # end - - if self._dict_has_key_from_group(self.ctx, info_groups["geometry_info"]): - output += "\n├─ Geometry info:" - if "geometry_type" in self.ctx.keys(): - printed_keys.append("geometry_type") - output += f"\n│ ├─ Type: {gkenums.gkyl_geometry_id[self.ctx['geometry_type']]:s}" - # end - if "geqdsk_file" in self.ctx.keys(): - printed_keys.append("geqdsk_file") - output += f"\n│ ├─ GEQDSK file: {self.ctx['geqdsk_file']:s}" - # end - if "geqdsk_sign_convention" in self.ctx.keys(): - printed_keys.append("geqdsk_sign_convention") - output += f"\n│ ├─ GEQDSK sign convention: {self.ctx['geqdsk_sign_convention']:d}" - # end - if "is_multib" in self.ctx.keys(): - printed_keys.append("is_multib") - if self.ctx['is_multib'] == 1: - output += f"\n│ ├─ Multiblock: yes" - else: - output += f"\n│ ├─ Multiblock: no" - # end - # end - if "topo_file" in self.ctx.keys(): - printed_keys.append("topo_file") - output += f"\n│ ├─ Block topology file: {self.ctx['topo_file']:s}" - # end - if "half_domain" in self.ctx.keys(): - printed_keys.append("half_domain") - if self.ctx['half_domain'] == 1: - output += f"\n│ ├─ Half domain: yes" - else: - output += f"\n│ ├─ Half domain: no" - # end - # end - # end - - # Print any other keys in the context that were not printed above - for key, val in self.ctx.items(): - if key not in sum(info_groups.values(), []): - output += f"\n├─ {key:s}: {val}" - # end - # end - - if self._dict_has_key_from_group(self.ctx, info_groups["species_info"]): - output += "\n├─ Species properties:" - if "mass" in self.ctx.keys(): - printed_keys.append("mass") - output += f"\n│ ├─ Mass: {self.ctx['mass']:e}" - # end - if "charge" in self.ctx.keys(): - printed_keys.append("charge") - output += f"\n│ ├─ Charge: {self.ctx['charge']:e}" - # end - if "gas_gamma" in self.ctx.keys(): - printed_keys.append("gas_gamma") - output += f"\n│ ├─ Adiabatic index: {self.ctx['gas_gamma']:e}" - # end - if "vdim" in self.ctx.keys(): - printed_keys.append("vdim") - output += f"\n│ ├─ Velocity dimensions: {self.ctx['vdim']:d}" - # end - # end - - return output - - # ---- Write ---- - def write(self, out_name: str = "", - extension: Literal["gkyl", "bp", "txt", "npy"] = "gkyl", - mode: str = "", var_name: str = "", append: bool = False, - cleaning: bool = True) -> None: - """Writes data in a file. - - The available formats are Gkeyll .gkyl (default), ADIOS .bp file, ASCII .txt file, - or NumPy .npy file. - - Args: - out_name: str - Specify output file name. - extension: str = "gkyl" - Specify file extension (extension). - var_name: str - Specify variable name for Adios. - append: bool = False - Allows for writing multiple datasets into one file. - cleaning: bool = True - Remove temporary files after writing. - - Returns: - None - """ - - if mode: - extension = mode - print("Deprecation warning: mode of the write method is going to be renamed to extension.") - # end - - if not out_name: - if self._file_name is not None: - fn = self._file_name - out_name = f"{fn.split('.', maxsplit=1)[0].strip('_')}_mod.{extension}" - else: - out_name = f"gdata.{extension}" - # end - else: - if not isinstance(out_name, str): - raise TypeError("'out_name' must be a string") - # end - if out_name.split(".")[-1] != extension: - out_name += "." + extension - # end - # end - - num_dims = self.num_dims - num_comps = self.num_comps - num_cells = self.num_cells - lo, up = self.bounds - values = self.values - - full_shape = list(num_cells) + [num_comps] - offset = [0] * (num_dims + 1) - - if not var_name: - var_name = self._var_name - # end - - if extension == "bp": - if not has_adios: - raise ModuleNotFoundError("ADIOS2 is not installed") - # end - - if not append: - fh = adios2.open(out_name, "w", engine_type="BP3") - fh.write_attribute("numCells", num_cells) - fh.write_attribute("lowerBounds", lo) - fh.write_attribute("upperBounds", up) - - if self.ctx["time"]: - fh.write("time", self.ctx["time"]) - # end - else: - fh = adios2.open(out_name, "a", engine_type="BP3") - # end - fh.write(var_name, values, full_shape, offset, full_shape) - fh.close() - - if cleaning: - if len(out_name.split("/")) > 1: - nm = out_name.split("/")[-1] - else: - nm = out_name - # end - shutil.move(f"{out_name}.dir/{nm}.0", f"{out_name}") - shutil.rmtree(f"{out_name}.dir") - # end - elif extension == "gkyl": - dti = np.dtype("i8") - dtf = np.dtype("f8") - - fh = open(out_name, "w", encoding="utf-8") - - # sep='' results in a binary file - np.array([103, 107, 121, 108, 48], dtype=np.dtype("b")).tofile(fh, sep="") - # version 1 - np.array([1], dtype=dti).tofile(fh, sep="") - # type 1 - np.array([1], dtype=dti).tofile(fh, sep="") - # meta size - np.array([0], dtype=dti).tofile(fh, sep="") - # real type (double) - np.array([2], dtype=dti).tofile(fh, sep="") - # num dims - np.array([num_dims], dtype=dti).tofile(fh, sep="") - # num cells - np.array(num_cells, dtype=dti).tofile(fh, sep="") - # lower - np.array(lo, dtype=dtf).tofile(fh, sep="") - # upper - np.array(up, dtype=dtf).tofile(fh, sep="") - # elem_sz - np.array([num_comps * 8], dtype=dti).tofile(fh, sep="") - # asize - np.array([np.size(values)], dtype=dti).tofile(fh, sep="") - # data - np.array(values, dtype=dtf).tofile(fh, sep="") - - fh.close() - elif extension == "txt": - num_rows = np.prod(num_cells) - grid = self.get_grid() - for d in range(num_dims): - grid[d] = 0.5 * (grid[d][1:] + grid[d][:-1]) - # end - - basis = np.full(num_dims, 1.0) - for d in range(num_dims - 1): - basis[d] = np.prod(num_cells[(d + 1) :]) - # end - - fh = open(out_name, "w", encoding="utf-8") - for i in range(num_rows): - idx = i - idxs = np.zeros(num_dims, np.int32) - for d in range(num_dims): - idxs[d] = int(idx // basis[d]) - idx = idx % basis[d] - # end - line = "" - for d in range(num_dims): - line += f"{grid[d][idxs[d]]:.15e}, " - # end - for c in range(num_comps - 1): - line += f"{values[tuple(idxs)][c]:.15e}, " - # end - line += f"{values[tuple(idxs)][num_comps - 1]:.15e}\n" - fh.write(line) - # end - fh.close() - elif extension == "npy": - np.save(out_name, values.squeeze()) - # end - - # ---- Context (metadata) ---- - def get_ctx(self) -> dict: - return self.ctx - diff --git a/src/postgkyl/data/gkyl_adios_reader.py b/src/postgkyl/data/gkyl_adios_reader.py deleted file mode 100644 index fdce8bd9..00000000 --- a/src/postgkyl/data/gkyl_adios_reader.py +++ /dev/null @@ -1,341 +0,0 @@ -"""Module including Gkeyll ADIOS reader class.""" - -from typing import Tuple -import click -import numpy as np -import re - -try: - import adios2 - has_adios = True -except ModuleNotFoundError: - has_adios = False -# end - -import postgkyl.data.idx_parser as idx_parser - - -class GkylAdiosReader(object): - """Provides a framework to read gkyl ADIOS output.""" - - def __init__(self, file_name: str, ctx: dict | None = None, - var_name: str = "CartGridField", c2p: str = "", - axes: tuple | None = (None, None, None, None, None, None), - comp: int | slice | None = None, click_mode: bool = False, - **kwargs): - """Initialize the instance of ADIOS reader. - - Args: - file_name: str - ctx: dict - Passes context variable with metadata. - var_name: str = "CartGridField" - c2p: str - Allows to specify a name of the file containing c2p mapping. - axes: tuple - Coordinate indices for partial loading. - comp: int - Component index for partial loading. - click_mode: bool = False - Enables command-line behavior like prompting when a - var_name is either missing or doesn't match any available. - **kwargs - This is not directly used but allowes for unified interface to all the readers - we use. - """ - self._file_name = file_name - self.var_name = var_name - self.c2p = c2p - - self.axes = axes - self.comp = comp - - self.lower = None - self.upper = None - self.num_comps = None - self.cells = None - - self.is_frame = False - self.is_diagnostic = False - self.click_mode = click_mode - - self.ctx = ctx - if not ("grid_type" in self.ctx.keys()): - self.ctx["grid_type"] = "uniform" - - def is_compatible(self) -> bool: - """Checks if file can be read with Gkeyll ADIOS reader.""" - if not has_adios: - return False - # end - try: - fh = adios2.open(self._file_name, "rra") - for vn in fh.available_variables(): - if "TimeMesh" in vn: - self.is_diagnostic = True - fh.close() - return True - # end - # end - - available_var_names = "" - for vn in fh.available_variables(): - available_var_names += f"'{str(vn):s}', " - # end - if self.var_name not in fh.available_variables(): - self.ctx["var_names"] = available_var_names[:-2] - # end - self.is_frame = True - fh.close() - return True - except ModuleNotFoundError: - return False - except TypeError: - return False - # end - - def _create_offset_count(self, num_elems: np.ndarray, zs: tuple, comp: int | slice, - grid: list | None = None) -> Tuple[np.ndarray, np.ndarray]: - num_dims = len(num_elems) - count = np.copy(num_elems) - offset = np.zeros(num_dims, np.int32) - cnt = 0 - for d, z in enumerate(zs): - if d < num_dims - 1 and z is not None: # Last dim stores comp - z = idx_parser.idx_parser(z, grid[d]) - if isinstance(z, int): - offset[d] = z - count[d] = 1 - elif isinstance(z, slice): - offset[d] = z.start - count[d] = z.stop - z.start - else: - raise TypeError("'z' is neither number or slice") - # end - cnt = cnt + 1 - # end - # end - - if comp is not None: - comp = idx_parser.idx_parser(comp) - if isinstance(comp, int): - offset[-1] = comp - count[-1] = 1 - elif isinstance(comp, slice): - offset[-1] = comp.start - count[-1] = comp.stop - comp.start - else: - raise TypeError("'comp' is neither number or slice") - # end - cnt = cnt + 1 - # end - - if cnt > 0: - return tuple(offset), tuple(count) - else: - return (), () - # end - - def _preload_frame(self) -> None: - fh = adios2.open(self._file_name, "rra") - - # Postgkyl conventions require the attributes to be - # narrays even for 1D data - self.lower = np.atleast_1d(fh.read_attribute("lowerBounds")) - self.upper = np.atleast_1d(fh.read_attribute("upperBounds")) - self.cells = np.atleast_1d(fh.read_attribute("numCells")) - if "changeset" in fh.available_attributes().keys(): - self.ctx["changeset"] = fh.read_attribute_string("changeset")[0] - # end - if "builddate" in fh.available_attributes().keys(): - self.ctx["builddate"] = fh.read_attribute_string("builddate")[0] - # end - if "polyOrder" in fh.available_attributes().keys(): - self.ctx["poly_order"] = fh.read_attribute("polyOrder")[0] - self.ctx["is_modal"] = True - # end - if "basisType" in fh.available_attributes().keys(): - self.ctx["basis_type"] = fh.read_attribute_string("basisType")[0] - self.ctx["is_modal"] = True - # end - if "charge" in fh.available_attributes().keys(): - self.ctx["charge"] = fh.read_attribute("charge")[0] - # end - if "mass" in fh.available_attributes().keys(): - self.ctx["mass"] = fh.read_attribute("mass")[0] - # end - if "time" in fh.available_variables(): - self.ctx["time"] = fh.read("time") - # end - if "frame" in fh.available_variables(): - self.ctx["frame"] = fh.read("frame") - # end - - fh.close() - - def _load_frame(self) -> Tuple[list, np.ndarray]: - fh = adios2.open(self._file_name, "rra") - - if self.var_name not in fh.available_variables(): - if self.click_mode: - var_name = self.var_name - while True: - var_name = click.prompt(f"Variable name '{var_name:s}' is not available, please select from the available ones: {self.ctx['var_names']:s}") - if var_name in fh.available_variables(): - self.var_name = var_name - self.ctx.pop("var_names", None) - break - # end - # end - else: - raise ValueError( - f"Could not find the variable '{var_name:s}'; available variables are: {self.ctx['var_names']:s}" - ) - # end - # end - - num_dims = len(self.cells) - grid = [np.linspace(self.lower[d], self.upper[d], self.cells[d] + 1) for d in range(num_dims)] - var_shape = fh.available_variables()[self.var_name]["Shape"] - num_elems = np.array([v for v in var_shape.split(",")], dtype=np.int32) - offset, count = self._create_offset_count(num_elems, self.axes, self.comp, grid) - data = fh.read(self.var_name, start=offset, count=count) - - # Adjust boundaries for 'offset' and 'count' - dz = (self.upper - self.lower) / self.cells - if offset: - if self.ctx["grid_type"] == "uniform": - self.lower = self.lower + offset[:num_dims] * dz - self.cells = self.cells - offset[:num_dims] - elif self.ctx["grid_type"] == "mapped": - idx = np.full(num_dims, 0) - for d in range(num_dims): - self.lower[d] = self._grid[d][tuple(idx)] - self.cells[d] = self.cells[d] - offset[d] - # end - # end - # end - if count: - if self.ctx["grid_type"] == "uniform": - self.upper = self.lower + count[:num_dims] * dz - self.cells = count[:num_dims] - elif self.ctx["grid_type"] == "mapped": - idx = np.full(num_dims, 0) - for d in range(num_dims): - idx[-d - 1] = ( - count[d] - 1 - ) # .Reverse indexing of idx because of transpose() in composing self._grid. - self.upper[d] = self._grid[d][tuple(idx)] - self.cells[d] = count[d] - # end - # end - # end - - # Check for mapped grid ... - if self.c2p: - grid_fh = adios2.open(self.c2p, "rra") - grid_dims = grid_fh.available_variables()["CartGridField"]["Shape"] - grid_dims = [int(v) for v in grid_dims.split(",")] - offset, count = self._create_offset_count(grid_dims, self.axes, None) - tmp = grid_fh.read("CartGridField", start=offset, count=count) - num_comps = tmp.shape[-1] - num_coeff = num_comps / num_dims - grid = [ - tmp[..., int(d * num_coeff) : int((d + 1) * num_coeff)] - for d in range(num_dims) - ] - if self.ctx: - self.ctx["grid_type"] = "c2p" - # end - else: - # Create sparse unifrom grid - # Adjust for ghost cells - dz = (self.upper - self.lower) / self.cells - for d in range(num_dims): - if self.cells[d] != data.shape[d]: - ngl = int(np.floor((self.cells[d] - data.shape[d]) * 0.5)) - ngu = int(np.ceil((self.cells[d] - data.shape[d]) * 0.5)) - self.cells[d] = data.shape[d] - self.lower[d] = self.lower[d] - ngl * dz[d] - self.upper[d] = self.upper[d] + ngu * dz[d] - # end - # end - grid = [ - np.linspace(self.lower[d], self.upper[d], self.cells[d] + 1) - for d in range(num_dims) - ] - if self.ctx: - self.ctx["grid_type"] = "uniform" - # end - # end - - fh.close() - return grid, data - - def _load_diagnostic(self) -> Tuple[list, np.ndarray]: - - fh = adios2.open(self._file_name, "rra") - - def natural_sort(l): - convert = lambda text: int(text) if text.isdigit() else text.lower() - alphanum_key = lambda key: [convert(c) for c in re.split("([0-9]+)", key)] - return sorted(l, key=alphanum_key) - - time_lst = [key for key in fh.available_variables() if "TimeMesh" in key] - data_lst = [key for key in fh.available_variables() if "Data" in key] - time_lst = natural_sort(time_lst) - data_lst = natural_sort(data_lst) - - for i in range(len(data_lst)): - if i == 0: - data = np.atleast_1d(fh.read(data_lst[i])) - grid = np.atleast_1d(fh.read(time_lst[i])) - else: - next_data = np.atleast_1d(fh.read(data_lst[i])) - next_grid = np.atleast_1d(fh.read(time_lst[i])) - # deal with weird behavior after restart where some data - # doesn't have second dimension - if len(next_data.shape) < 2: - next_data = np.expand_dims(next_data, axis=1) - # end - data = np.append(data, next_data, axis=0) - grid = np.append(grid, next_grid, axis=0) - # end - # end - fh.close() - # end - - return [np.squeeze(grid)], data - - def preload(self) -> None: - """Loads metadata.""" - if self.is_frame: - self._preload_frame() - if self.ctx: - self.ctx["cells"] = self.cells - self.ctx["lower"] = self.lower - self.ctx["upper"] = self.upper - # end - # end - - def load(self) -> Tuple[list, np.ndarray]: - """Loads data. - - Returns: - A tuple including a grid list and a data NumPy array - - Notes: - Needs to be called after the preload. - """ - grid, data = None, None - - if self.is_frame: - grid, data = self._load_frame() - # end - if self.is_diagnostic: - grid, data = self._load_diagnostic() - # end - - self.ctx["num_comps"] = data.shape[-1] - - return grid, data diff --git a/src/postgkyl/data/gkyl_h5_reader.py b/src/postgkyl/data/gkyl_h5_reader.py deleted file mode 100644 index 0da26a6d..00000000 --- a/src/postgkyl/data/gkyl_h5_reader.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Module including legacy Gkeyll reader class""" - -from typing import Tuple -import numpy as np -import tables - - -class GkylH5Reader(object): - """Provides a framework to read legacy Gkeyll HDF5 output""" - - def __init__(self, file_name: str, ctx: dict | None = None, **kwargs): - """Initialize the instance of Gkeyll reader. - - Args: - file_name: str - ctx: dict - Passes context variable with metadata. - **kwargs - This is not directly used but allowes for unified interface to all the readers - we use. - """ - self._file_name = file_name - - self.is_frame = False - self.is_diagnostic = False - - self.ctx = ctx - - def is_compatible(self) -> bool: - """Checks if file can be read with the legacy Gkeyll HDF5 reader.""" - try: - fh = tables.open_file(self._file_name, "r") - - if "/DataStruct/data" in fh: - self.is_diagnostic = True - # end - if "/StructGridField" in fh: - self.is_frame = True - # end - - fh.close() - except: - return False - # end - return self.is_frame or self.is_diagnostic - - def _read_frame(self) -> tuple: - fh = tables.open_file(self._file_name, "r") - - # Postgkyl conventions require the attributes to be - # narrays even for 1D data - lower = np.atleast_1d(fh.root.StructGrid._v_attrs.vsLowerBounds) - upper = np.atleast_1d(fh.root.StructGrid._v_attrs.vsUpperBounds) - cells = np.atleast_1d(fh.root.StructGrid._v_attrs.vsNumCells) - if "/timeData" in fh: - self.ctx["time"] = fh.root.timeData._v_attrs.vsTime - # end - - data = fh.root.StructGridField.read() - - fh.close() - return cells, lower, upper, data - - def _read_diagnostic(self): - fh = tables.open_file(self._file_name, "r") - - grid = fh.root.DataStruct.timeMesh.read() - data = fh.root.DataStruct.data.read() - - fh.close() - # end - - return [np.squeeze(grid)], [grid[0]], [grid[-1]], data - - def preload(self) -> None: - """Loads metadata.""" - pass - - def load(self) -> Tuple[list, np.ndarray]: - """Loads data. - - Returns: - A tuple including a grid list and a data NumPy array - - Notes: - Needs to be called after the preload. - """ - grid = None - - if self.is_frame: - cells, lower, upper, data = self._read_frame() - else: - grid, lower, upper, data = self._read_diagnostic() - cells = grid[0].shape - # end - - if self.ctx: - self.ctx["cells"] = cells - self.ctx["lower"] = lower - self.ctx["upper"] = upper - self.ctx["num_comps"] = 1 - if len(data.shape) > len(cells): - self.ctx["num_comps"] = data.shape[-1] - # end - # end - - num_dims = len(cells) - grid = [np.linspace(lower[d], upper[d], cells[d] + 1) for d in range(num_dims)] - if self.ctx: - self.ctx["grid_type"] = "uniform" - # end - - return grid, data diff --git a/src/postgkyl/data/idx_parser.py b/src/postgkyl/data/idx_parser.py deleted file mode 100644 index 2754f360..00000000 --- a/src/postgkyl/data/idx_parser.py +++ /dev/null @@ -1,102 +0,0 @@ -import numpy as np - -def _find_nearest_index(array, value): - if array is None: - raise TypeError("The index value is float but the 'array' from which to select the neares value is not specified.") - # end - idx = np.searchsorted(array, value) - if idx == len(array): - return int(idx - 2) - elif idx > 0: - return int(idx - 1) - else: - return int(idx) - # end - - -def _find_cell_index(array, value): - if array is None: - raise TypeError("The index value is float but the 'array' from which to select the neares value is not specified.") - # end - idx = np.searchsorted(array, value) - return int(idx) - - -def _is_int_str(value: str) -> bool: - """Whether the string represents an integer (e.g. '2' or '-1'), not a float.""" - try: - int(value) - return True - except ValueError: - return False - # end - - -def _resolve_negative_index(idx: int, array: np.ndarray | None, nodal: bool) -> int: - """Translate a Python-style negative index into a positive one. - - '-1' refers to the last cell, '-2' to the one before, etc. The number - of cells is the length of the grid array for nodal data and one less - for cell-centered data (where the grid stores cell edges). - """ - if idx < 0 and array is not None: - num_cells = len(array) if nodal else len(array) - 1 - idx += num_cells - # end - return idx - - -def _string_to_index(value: str, array: np.ndarray, nodal: bool = False) -> int: - if isinstance(value, str): - if _is_int_str(value): - return _resolve_negative_index(int(value), array, nodal) - else: - if nodal: - return _find_cell_index(array, float(value)) - else: - return _find_nearest_index(array, float(value)) - # end - # end - else: - raise TypeError("Value is not string") - # end - - -def idx_parser(value: int | float | str, array: np.ndarray | None = None, - nodal: bool = False) -> int | slice: - idx = None - if isinstance(value, int): - idx = _resolve_negative_index(value, array, nodal) - elif isinstance(value, float): - if nodal: - idx = _find_cell_index(array, value) - else: - idx = _find_nearest_index(array, value) - # end - else: - if isinstance(value, str): - if len(value.split(",")) > 1: - idxs = value.split(",") - idx = tuple([_string_to_index(i, array, nodal) for i in idxs]) - elif len(value.split(":")) == 2: - idxs = value.split(":") - if idxs[0] == "": - idxs[0] = str(0) - # end - if idxs[1] == "": - idxs[1] = str(len(array)) - # end - try: - if int(idxs[1]) < 0: - idxs[1] = str(len(array) + int(idxs[1]) + 1) - # end - except ValueError: - pass - idx = slice(_string_to_index(idxs[0], array, nodal), _string_to_index(idxs[1], array, nodal)) - else: - idx = _string_to_index(value, array, nodal) - # end - # end - # end - - return idx diff --git a/src/postgkyl/data/old/recovData.py b/src/postgkyl/data/old/recovData.py deleted file mode 100644 index fa0cbe3d..00000000 --- a/src/postgkyl/data/old/recovData.py +++ /dev/null @@ -1,445 +0,0 @@ -# --------------------------------------------------------------------- -# -- P1 --------------------------------------------------------------- - - -def p1e(x, fL, fR, dx): - return ( - (3.061862178478972 * fR[1] * x**3) / dx**3 - + (3.061862178478972 * fL[1] * x**3) / dx**3 - - (1.767766952966368 * fR[0] * x**3) / dx**3 - + (1.767766952966368 * fL[0] * x**3) / dx**3 - + (1.224744871391589 * fR[1] * x**2) / dx**2 - - (1.224744871391589 * fL[1] * x**2) / dx**2 - - (1.530931089239486 * fR[1] * x) / dx - - (1.530931089239486 * fL[1] * x) / dx - + (1.590990257669731 * fR[0] * x) / dx - - (1.590990257669731 * fL[0] * x) / dx - - 0.408248290463863 * fR[1] - + 0.408248290463863 * fL[1] - + 0.3535533905932737 * fR[0] - + 0.3535533905932737 * fL[0] - ) - - -def p1c1(x, f, fL, fR, dx): - return ( - (23.57633877428808 * fR[1] * x**5) / dx**5 - + (23.57633877428808 * fL[1] * x**5) / dx**5 - + (81.44553394754065 * f[1] * x**5) / dx**5 - - (18.56155300614687 * fR[0] * x**5) / dx**5 - + (18.56155300614687 * fL[0] * x**5) / dx**5 - + (2.296396633859228 * fR[1] * x**4) / dx**4 - - (2.296396633859228 * fL[1] * x**4) / dx**4 - - (1.325825214724776 * fR[0] * x**4) / dx**4 - - (1.325825214724776 * fL[0] * x**4) / dx**4 - + (2.651650429449552 * f[0] * x**4) / dx**4 - - (12.5026038954558 * fR[1] * x**3) / dx**3 - - (12.5026038954558 * fL[1] * x**3) / dx**3 - - (45.41762231410475 * f[1] * x**3) / dx**3 - + (10.16465997955662 * fR[0] * x**3) / dx**3 - - (10.16465997955662 * fL[0] * x**3) / dx**3 - - (1.913663861549357 * fR[1] * x**2) / dx**2 - + (1.913663861549357 * fL[1] * x**2) / dx**2 - + (1.458407736197253 * fR[0] * x**2) / dx**2 - + (1.458407736197253 * fL[0] * x**2) / dx**2 - - (2.916815472394507 * f[0] * x**2) / dx**2 - + (1.243881510007081 * fR[1] * x) / dx - + (1.243881510007081 * fL[1] * x) / dx - + (7.08055628773262 * f[1] * x) / dx - - (1.027514541411701 * fR[0] * x) / dx - + (1.027514541411701 * fL[0] * x) / dx - + 0.130767030539206 * fR[1] - - 0.130767030539206 * fL[1] - - 0.104961162832378 * fR[0] - - 0.104961162832378 * fL[0] - + 0.917029106851303 * f[0] - ) - - -def p1c0(x, f, fL, fR, dx): - return ( - (-(4.082482904638631 * fR[1] * x**3) / dx**3) - - (4.082482904638631 * fL[1] * x**3) / dx**3 - - (16.32993161855453 * f[1] * x**3) / dx**3 - + (3.535533905932737 * fR[0] * x**3) / dx**3 - - (3.535533905932737 * fL[0] * x**3) / dx**3 - - (1.224744871391589 * fR[1] * x**2) / dx**2 - + (1.224744871391589 * fL[1] * x**2) / dx**2 - + (1.060660171779821 * fR[0] * x**2) / dx**2 - + (1.060660171779821 * fL[0] * x**2) / dx**2 - - (2.121320343559642 * f[0] * x**2) / dx**2 - + (0.6123724356957944 * fR[1] * x) / dx - + (0.6123724356957944 * fL[1] * x) / dx - + (4.898979485566357 * f[1] * x) / dx - - (0.5303300858899105 * fR[0] * x) / dx - + (0.5303300858899105 * fL[0] * x) / dx - + 0.1020620726159657 * fR[1] - - 0.1020620726159657 * fL[1] - - 0.0883883476483184 * fR[0] - - 0.0883883476483184 * fL[0] - + 0.883883476483184 * f[0] - ) - - -# --------------------------------------------------------------------- -# -- P2 --------------------------------------------------------------- - - -def p2e(x, fL, fR, dx): - return ( - (13.28156617270719 * fR[2] * x**5) / dx**5 - - (13.28156617270719 * fL[2] * x**5) / dx**5 - - (12.85982114961168 * fR[1] * x**5) / dx**5 - - (12.85982114961168 * fL[1] * x**5) / dx**5 - + (7.424621202458747 * fR[0] * x**5) / dx**5 - - (7.424621202458747 * fL[0] * x**5) / dx**5 - + (5.188111786213744 * fR[2] * x**4) / dx**4 - + (5.188111786213744 * fL[2] * x**4) / dx**4 - - (1.339564703084549 * fR[1] * x**4) / dx**4 - + (1.339564703084549 * fL[1] * x**4) / dx**4 - - (12.64911064067352 * fR[2] * x**3) / dx**3 - + (12.64911064067352 * fL[2] * x**3) / dx**3 - + (15.30931089239486 * fR[1] * x**3) / dx**3 - + (15.30931089239486 * fL[1] * x**3) / dx**3 - - (8.838834764831843 * fR[0] * x**3) / dx**3 - + (8.838834764831843 * fL[0] * x**3) / dx**3 - - (4.150489428970996 * fR[2] * x**2) / dx**2 - - (4.150489428970996 * fL[2] * x**2) / dx**2 - + (2.296396633859228 * fR[1] * x**2) / dx**2 - - (2.296396633859228 * fL[1] * x**2) / dx**2 - + (1.897366596101028 * fR[2] * x) / dx - - (1.897366596101028 * fL[2] * x) / dx - - (3.368048396326869 * fR[1] * x) / dx - - (3.368048396326869 * fL[1] * x) / dx - + (2.651650429449552 * fR[0] * x) / dx - - (2.651650429449552 * fL[0] * x) / dx - + 0.3458741190809163 * fR[2] - + 0.3458741190809163 * fL[2] - - 0.4975526040028326 * fR[1] - + 0.4975526040028326 * fL[1] - + 0.3535533905932737 * fR[0] - + 0.3535533905932737 * fL[0] - ) - - -def p2c1(x, f, fL, fR, dx): - return ( - (-(105.4224314958633 * fR[2] * x**6) / dx**6) - - (105.4224314958633 * fL[2] * x**6) / dx**6 - + (559.4859750252903 * f[2] * x**6) / dx**6 - + (138.2430773583255 * fR[1] * x**6) / dx**6 - - (138.2430773583255 * fL[1] * x**6) / dx**6 - - (92.80776503073433 * fR[0] * x**6) / dx**6 - - (92.80776503073433 * fL[0] * x**6) / dx**6 - + (185.6155300614687 * f[0] * x**6) / dx**6 - - (15.77185983008978 * fR[2] * x**5) / dx**5 - + (15.77185983008978 * fL[2] * x**5) / dx**5 - + (18.21807996194988 * fR[1] * x**5) / dx**5 - + (18.21807996194988 * fL[1] * x**5) / dx**5 - + (40.72276697377032 * f[1] * x**5) / dx**5 - - (11.13693180368812 * fR[0] * x**5) / dx**5 - + (11.13693180368812 * fL[0] * x**5) / dx**5 - + (56.03160729110844 * fR[2] * x**4) / dx**4 - + (56.03160729110844 * fL[2] * x**4) / dx**4 - - (319.5876860307667 * f[2] * x**4) / dx**4 - - (75.0156233727348 * fR[1] * x**4) / dx**4 - + (75.0156233727348 * fL[1] * x**4) / dx**4 - + (51.04427076690387 * fR[0] * x**4) / dx**4 - + (51.04427076690387 * fL[0] * x**4) / dx**4 - - (102.0885415338078 * f[0] * x**4) / dx**4 - + (9.091548272984086 * fR[2] * x**3) / dx**3 - - (9.091548272984086 * fL[2] * x**3) / dx**3 - - (11.48198316929614 * fR[1] * x**3) / dx**3 - - (11.48198316929614 * fL[1] * x**3) / dx**3 - - (29.08769069555023 * f[1] * x**3) / dx**3 - + (7.513009550107064 * fR[0] * x**3) / dx**3 - - (7.513009550107064 * fL[0] * x**3) / dx**3 - - (7.300414442029338 * fR[2] * x**2) / dx**2 - - (7.300414442029338 * fL[2] * x**2) / dx**2 - + (52.99285610204038 * f[2] * x**2) / dx**2 - + (9.903210483517913 * fR[1] * x**2) / dx**2 - - (9.903210483517913 * fL[1] * x**2) / dx**2 - - (6.794854225464475 * fR[0] * x**2) / dx**2 - - (6.794854225464475 * fL[0] * x**2) / dx**2 - + (13.58970845092895 * f[0] * x**2) / dx**2 - - (0.9412717097844931 * fR[2] * x) / dx - + (0.9412717097844931 * fL[2] * x) / dx - + (1.234313190699334 * fR[1] * x) / dx - + (1.234313190699334 * fL[1] * x) / dx - + (5.721854946032574 * f[1] * x) / dx - - (0.8286407592029846 * fR[0] * x) / dx - + (0.8286407592029846 * fL[0] * x) / dx - + 0.1432907064763796 * fR[2] - + 0.1432907064763796 * fL[2] - - 1.670077889276424 * f[2] - - 0.1961505458088089 * fR[1] - + 0.1961505458088089 * fL[1] - + 0.1353446573364875 * fR[0] - + 0.1353446573364875 * fL[0] - + 0.4364174665135718 * f[0] - ) - - -def p2c0(x, f, fL, fR, dx): - return ( - (12.10559416783207 * fR[2] * x**4) / dx**4 - + (12.10559416783207 * fL[2] * x**4) / dx**4 - - (86.4685297702291 * f[2] * x**4) / dx**4 - - (17.41434114009914 * fR[1] * x**4) / dx**4 - + (17.41434114009914 * fL[1] * x**4) / dx**4 - + (12.37436867076458 * fR[0] * x**4) / dx**4 - + (12.37436867076458 * fL[0] * x**4) / dx**4 - - (24.74873734152916 * f[0] * x**4) / dx**4 - + (3.458741190809164 * fR[2] * x**3) / dx**3 - - (3.458741190809164 * fL[2] * x**3) / dx**3 - - (4.975526040028328 * fR[1] * x**3) / dx**3 - - (4.975526040028328 * fL[1] * x**3) / dx**3 - - (14.54384534777511 * f[1] * x**3) / dx**3 - + (3.535533905932737 * fR[0] * x**3) / dx**3 - - (3.535533905932737 * fL[0] * x**3) / dx**3 - - (2.594055893106872 * fR[2] * x**2) / dx**2 - - (2.594055893106872 * fL[2] * x**2) / dx**2 - + (28.01580364555422 * f[2] * x**2) / dx**2 - + (3.731644530021244 * fR[1] * x**2) / dx**2 - - (3.731644530021244 * fL[1] * x**2) / dx**2 - - (2.651650429449552 * fR[0] * x**2) / dx**2 - - (2.651650429449552 * fL[0] * x**2) / dx**2 - + (5.303300858899105 * f[0] * x**2) / dx**2 - - (0.5188111786213743 * fR[2] * x) / dx - + (0.5188111786213743 * fL[2] * x) / dx - + (0.7463289060042488 * fR[1] * x) / dx - + (0.7463289060042488 * fL[1] * x) / dx - + (4.631066544949443 * f[1] * x) / dx - - (0.5303300858899105 * fR[0] * x) / dx - + (0.5303300858899105 * fL[0] * x) / dx - + 0.06485139732767176 * fR[2] - + 0.06485139732767176 * fL[2] - - 1.253793681668321 * f[2] - - 0.09329111325053105 * fR[1] - + 0.09329111325053105 * fL[1] - + 0.06629126073623878 * fR[0] - + 0.06629126073623878 * fL[0] - + 0.5745242597140695 * f[0] - ) - - -# --------------------------------------------------------------------- -# -- P3 --------------------------------------------------------------- - - -def p3e(x, fL, fR, dx): - return ( - (57.87876270165938 * fR[3] * x**7) / dx**7 - + (57.87876270165938 * fL[3] * x**7) / dx**7 - - (75.0052732521187 * fR[2] * x**7) / dx**7 - + (75.0052732521187 * fL[2] * x**7) / dx**7 - + (63.15090743112876 * fR[1] * x**7) / dx**7 - + (63.15090743112876 * fL[1] * x**7) / dx**7 - - (36.46019340493134 * fR[0] * x**7) / dx**7 - + (36.46019340493134 * fL[0] * x**7) / dx**7 - + (22.44994432064365 * fR[3] * x**6) / dx**6 - - (22.44994432064365 * fL[3] * x**6) / dx**6 - - (12.45146828691299 * fR[2] * x**6) / dx**6 - - (12.45146828691299 * fL[2] * x**6) / dx**6 - + (3.214955287402919 * fR[1] * x**6) / dx**6 - - (3.214955287402919 * fL[1] * x**6) / dx**6 - - (81.03026778232312 * fR[3] * x**5) / dx**5 - - (81.03026778232312 * fL[3] * x**5) / dx**5 - + (118.2889487256734 * fR[2] * x**5) / dx**5 - - (118.2889487256734 * fL[2] * x**5) / dx**5 - - (101.2710915531919 * fR[1] * x**5) / dx**5 - - (101.2710915531919 * fL[1] * x**5) / dx**5 - + (58.46889196936262 * fR[0] * x**5) / dx**5 - - (58.46889196936262 * fL[0] * x**5) / dx**5 - - (28.06243040080456 * fR[3] * x**4) / dx**4 - + (28.06243040080456 * fL[3] * x**4) / dx**4 - + (20.75244714485499 * fR[2] * x**4) / dx**4 - + (20.75244714485499 * fL[2] * x**4) / dx**4 - - (5.3582588123382 * fR[1] * x**4) / dx**4 - + (5.3582588123382 * fL[1] * x**4) / dx**4 - + (28.93938135082968 * fR[3] * x**3) / dx**3 - + (28.93938135082968 * fL[3] * x**3) / dx**3 - - (50.15174726673287 * fR[2] * x**3) / dx**3 - + (50.15174726673287 * fL[2] * x**3) / dx**3 - + (46.88476460795923 * fR[1] * x**3) / dx**3 - + (46.88476460795923 * fL[1] * x**3) / dx**3 - - (27.0689314672975 * fR[0] * x**3) / dx**3 - + (27.0689314672975 * fL[0] * x**3) / dx**3 - + (8.017837257372731 * fR[3] * x**2) / dx**2 - - (8.017837257372731 * fL[3] * x**2) / dx**2 - - (8.597442388582778 * fR[2] * x**2) / dx**2 - - (8.597442388582778 * fL[2] * x**2) / dx**2 - + (3.444594950788841 * fR[1] * x**2) / dx**2 - - (3.444594950788841 * fL[1] * x**2) / dx**2 - - (1.929292090055312 * fR[3] * x) / dx - - (1.929292090055312 * fL[3] * x) / dx - + (4.397542371171649 * fR[2] * x) / dx - - (4.397542371171649 * fL[2] * x) / dx - - (5.473078644031159 * fR[1] * x) / dx - - (5.473078644031159 * fL[1] * x) / dx - + (3.866990209613929 * fR[0] * x) / dx - - (3.866990209613929 * fL[0] * x) / dx - - 0.2672612419124243 * fR[3] - + 0.2672612419124243 * fL[3] - + 0.4941058844013091 * fR[2] - + 0.4941058844013091 * fL[2] - - 0.5358258812338199 * fR[1] - + 0.5358258812338199 * fL[1] - + 0.3535533905932737 * fR[0] - + 0.3535533905932737 * fL[0] - ) - - -def p3c1(x, f, fL, fR, dx): - return ( - (401.8439810429493 * fR[3] * x**7) / dx**7 - + (401.8439810429493 * fL[3] * x**7) / dx**7 - + (3132.067901626939 * f[3] * x**7) / dx**7 - - (688.0918546172629 * fR[2] * x**7) / dx**7 - + (688.0918546172629 * fL[2] * x**7) / dx**7 - + (699.7120543369067 * fR[1] * x**7) / dx**7 - + (699.7120543369067 * fL[1] * x**7) / dx**7 - + (1682.34017396527 * f[1] * x**7) / dx**7 - - (444.8143595401623 * fR[0] * x**7) / dx**7 - + (444.8143595401623 * fL[0] * x**7) / dx**7 - + (71.73458771205661 * fR[3] * x**6) / dx**6 - - (71.73458771205661 * fL[3] * x**6) / dx**6 - - (115.1760816539451 * fR[2] * x**6) / dx**6 - - (115.1760816539451 * fL[2] * x**6) / dx**6 - + (329.9639096031941 * f[2] * x**6) / dx**6 - + (110.11221859355 * fR[1] * x**6) / dx**6 - - (110.11221859355 * fL[1] * x**6) / dx**6 - - (67.28562964728236 * fR[0] * x**6) / dx**6 - - (67.28562964728236 * fL[0] * x**6) / dx**6 - + (134.5712592945648 * f[0] * x**6) / dx**6 - - (225.4640892514639 * fR[3] * x**5) / dx**5 - - (225.4640892514639 * fL[3] * x**5) / dx**5 - - (1898.949587184442 * f[3] * x**5) / dx**5 - + (390.6648175018948 * fR[2] * x**5) / dx**5 - - (390.6648175018948 * fL[2] * x**5) / dx**5 - - (401.4675415144392 * fR[1] * x**5) / dx**5 - - (401.4675415144392 * fL[1] * x**5) / dx**5 - - (976.5426685486359 * f[1] * x**5) / dx**5 - + (256.8454897225571 * fR[0] * x**5) / dx**5 - - (256.8454897225571 * fL[0] * x**5) / dx**5 - - (39.24355501362508 * fR[3] * x**4) / dx**4 - + (39.24355501362508 * fL[3] * x**4) / dx**4 - + (65.28373997652294 * fR[2] * x**4) / dx**4 - + (65.28373997652294 * fL[2] * x**4) / dx**4 - - (213.5772685324658 * f[2] * x**4) / dx**4 - - (64.6339969238295 * fR[1] * x**4) / dx**4 - + (64.6339969238295 * fL[1] * x**4) / dx**4 - + (40.41004769046555 * fR[0] * x**4) / dx**4 - + (40.41004769046555 * fL[0] * x**4) / dx**4 - - (80.82009538093111 * f[0] * x**4) / dx**4 - + (35.99152857394851 * fR[3] * x**3) / dx**3 - + (35.99152857394851 * fL[3] * x**3) / dx**3 - + (357.2844328894097 * f[3] * x**3) / dx**3 - - (62.90585540784163 * fR[2] * x**3) / dx**3 - + (62.90585540784163 * fL[2] * x**3) / dx**3 - + (65.13633368748619 * fR[1] * x**3) / dx**3 - + (65.13633368748619 * fL[1] * x**3) / dx**3 - + (159.7430908428324 * f[1] * x**3) / dx**3 - - (41.86016901907076 * fR[0] * x**3) / dx**3 - + (41.86016901907076 * fL[0] * x**3) / dx**3 - + (5.206896265774277 * fR[3] * x**2) / dx**2 - - (5.206896265774277 * fL[3] * x**2) / dx**2 - - (8.847583492560934 * fR[2] * x**2) / dx**2 - - (8.847583492560934 * fL[2] * x**2) / dx**2 - + (40.52285884446233 * f[2] * x**2) / dx**2 - + (8.934418153608549 * fR[1] * x**2) / dx**2 - - (8.934418153608549 * fL[1] * x**2) / dx**2 - - (5.655473181560367 * fR[0] * x**2) / dx**2 - - (5.655473181560367 * fL[0] * x**2) / dx**2 - + (11.31094636312074 * f[0] * x**2) / dx**2 - - (1.45245001097914 * fR[3] * x) / dx - - (1.45245001097914 * fL[3] * x) / dx - - (19.04079750242088 * f[3] * x) / dx - + (2.555453870888018 * fR[2] * x) / dx - - (2.555453870888018 * fL[2] * x) / dx - - (2.661188807467072 * fR[1] * x) / dx - - (2.661188807467072 * fL[1] * x) / dx - - (4.116769382158051 * f[1] * x) / dx - + (1.71597690551618 * fR[0] * x) / dx - - (1.71597690551618 * fL[0] * x) / dx - - 0.1034854320490978 * fR[3] - + 0.1034854320490978 * fL[3] - + 0.1783413426510973 * fR[2] - + 0.1783413426510973 * fL[2] - - 1.443715630985074 * f[2] - - 0.1823960868039229 * fR[1] - + 0.1823960868039229 * fL[1] - + 0.116354973271419 * fR[0] - + 0.116354973271419 * fL[0] - + 0.4743968346437084 * f[0] - ) - - -def p3c0(x, f, fL, fR, dx): - return ( - (-(33.67491648096548 * fR[3] * x**5) / dx**5) - - (33.67491648096548 * fL[3] * x**5) / dx**5 - - (404.0989977715859 * f[3] * x**5) / dx**5 - + (62.25734143456496 * fR[2] * x**5) / dx**5 - - (62.25734143456496 * fL[2] * x**5) / dx**5 - - (67.5140610354613 * fR[1] * x**5) / dx**5 - - (67.5140610354613 * fL[1] * x**5) / dx**5 - - (173.6075855197576 * f[1] * x**5) / dx**5 - + (44.54772721475249 * fR[0] * x**5) / dx**5 - - (44.54772721475249 * fL[0] * x**5) / dx**5 - - (9.354143466934852 * fR[3] * x**4) / dx**4 - + (9.354143466934852 * fL[3] * x**4) / dx**4 - + (17.29370595404582 * fR[2] * x**4) / dx**4 - + (17.29370595404582 * fL[2] * x**4) / dx**4 - - (76.09230619780162 * f[2] * x**4) / dx**4 - - (18.75390584318369 * fR[1] * x**4) / dx**4 - + (18.75390584318369 * fL[1] * x**4) / dx**4 - + (12.37436867076458 * fR[0] * x**4) / dx**4 - + (12.37436867076458 * fL[0] * x**4) / dx**4 - - (24.74873734152916 * f[0] * x**4) / dx**4 - + (9.354143466934852 * fR[3] * x**3) / dx**3 - + (9.354143466934852 * fL[3] * x**3) / dx**3 - + (149.6662954709577 * f[3] * x**3) / dx**3 - - (17.29370595404582 * fR[2] * x**3) / dx**3 - + (17.29370595404582 * fL[2] * x**3) / dx**3 - + (18.75390584318369 * fR[1] * x**3) / dx**3 - + (18.75390584318369 * fL[1] * x**3) / dx**3 - + (48.22432931104378 * f[1] * x**3) / dx**3 - - (12.37436867076458 * fR[0] * x**3) / dx**3 - + (12.37436867076458 * fL[0] * x**3) / dx**3 - + (2.004459314343182 * fR[3] * x**2) / dx**2 - - (2.004459314343182 * fL[3] * x**2) / dx**2 - - (3.705794133009818 * fR[2] * x**2) / dx**2 - - (3.705794133009818 * fL[2] * x**2) / dx**2 - + (25.79232716574833 * f[2] * x**2) / dx**2 - + (4.018694109253648 * fR[1] * x**2) / dx**2 - - (4.018694109253648 * fL[1] * x**2) / dx**2 - - (2.651650429449552 * fR[0] * x**2) / dx**2 - - (2.651650429449552 * fL[0] * x**2) / dx**2 - + (5.303300858899105 * f[0] * x**2) / dx**2 - - (0.5011148285857955 * fR[3] * x) / dx - - (0.5011148285857955 * fL[3] * x) / dx - - (11.62586402319046 * f[3] * x) / dx - + (0.926448533252454 * fR[2] * x) / dx - - (0.926448533252454 * fL[2] * x) / dx - - (1.004673527313412 * fR[1] * x) / dx - - (1.004673527313412 * fL[1] * x) / dx - - (0.1339564703084549 * f[1] * x) / dx - + (0.6629126073623879 * fR[0] * x) / dx - - (0.6629126073623879 * fL[0] * x) / dx - - 0.05011148285857954 * fR[3] - + 0.05011148285857954 * fL[3] - + 0.0926448533252454 * fR[2] - + 0.0926448533252454 * fL[2] - - 1.198206769673174 * f[2] - - 0.1004673527313412 * fR[1] - + 0.1004673527313412 * fL[1] - + 0.06629126073623878 * fR[0] - + 0.06629126073623878 * fL[0] - + 0.5745242597140695 * f[0] - ) - - -recovC0Fn = [p1c0, p2c0, p3c0] -recovC1Fn = [p1c1, p2c1, p3c1] -recovEdFn = [p1e, p2e, p3e] diff --git a/src/postgkyl/data/old/three_cell_recov.mac b/src/postgkyl/data/old/three_cell_recov.mac deleted file mode 100644 index fb9f7df4..00000000 --- a/src/postgkyl/data/old/three_cell_recov.mac +++ /dev/null @@ -1,73 +0,0 @@ -kill(all) $ - -load("modal-basis")$ -load("basis-precalc/basisSer1x1v")$ -poly_order : 3 $ -basisX : basisC[poly_order] -N : length(basisX)$ -eta(xc, dx, basis) := subst(wx=x, subst(x=(wx-xc)/(dx/2), basis))$ -baL : eta(-dx/2, dx, basisX) $ -baR : eta(dx/2, dx, basisX) $ -baC : eta(0, dx, basisX) $ - -r1p : doExpand(r1, create_list(x^i, i, 0, 2*N-1)) $ -eqList1 : append( - calcInnerProdListGen([x], [[-dx,0]], 1, baL, r1p-doExpand(qL, baL)), - calcInnerProdListGen([x], [[0,dx]], 1, baR, r1p-doExpand(qR, baR)) -) $ -r1Sol : linsolve(eqList1, makelist(r1[i], i, 1, 2*N)) $ -r1s : fullratsimp(subst(r1Sol, r1p)) $ -substList : append( - makelist(qR[i]=fR[i-1],i,1,N), - makelist(qL[i]=fL[i-1],i,1,N) -) $ -out : float(expand(subst(substList, r1s))) $ -fh : openw("~/max-out/pgkyl_recov")$ -printf(fh, sconcat("def p", poly_order, "e(x, fL, fR, dx):~%"))$ -printf(fh, " return ~a~%~%", out) $ - -qh : sum(q[j,i]*baC[i], i, 1,N) $ -subListR : append( - makelist(qL[i]=q[j,i], i,1,N), makelist(qR[i]=q[j+1,i], i,1,N) -) $ -subListL : append( - makelist(qL[i]=q[j-1,i], i,1,N), makelist(qR[i]=q[j,i], i,1,N) -) $ -der : subst(x=0, diff(r1s, x)) $ -val : subst(x=0, r1s) $ -derL : subst(subListL, der) $ -derR : subst(subListR, der) $ -valL : subst(subListL, val) $ -valR : subst(subListR, val) $ - -r2p : doExpand(r2, create_list(x^i, i, 0, N-1+4)) $ -eqList2 : append( - [derL-subst(x=-dx/2, diff(r2p, x))], - [derR-subst(x=dx/2, diff(r2p, x))], - [valL-subst(x=-dx/2, r2p)], - [valR-subst(x=dx/2, r2p)], - calcInnerProdListGen([x], [[-dx/2,dx/2]], 1, baC, r2p-qh) -)$ -r2Sol : linsolve(eqList2, makelist(r2[i], i, 1, N+4)) $ -r2s : fullratsimp(subst(r2Sol, r2p)) $ -substList : append( - makelist(q[j+1,i]=fR[i-1],i,1,N), - makelist(q[j,i]=f[i-1],i,1,N), - makelist(q[j-1,i]=fL[i-1],i,1,N) -) $ -out : float(expand(subst(substList, r2s))) $ -printf(fh, sconcat("def p", poly_order, "c1(x, f, fL, fR, dx):~%"))$ -printf(fh, " return ~a~%~%", out) $ - -r2p : doExpand(r2, create_list(x^i, i, 0, N-1+2)) $ -eqList2 : append( - [valL-subst(x=-dx/2, r2p)], - [valR-subst(x=dx/2, r2p)], - calcInnerProdListGen([x], [[-dx/2,dx/2]], 1, baC, r2p-qh) -)$ -r2Sol : linsolve(eqList2, makelist(r2[i], i, 1, N+2)) $ -r2s : fullratsimp(subst(r2Sol, r2p)) $ -out : float(expand(subst(substList, r2s))) $ -printf(fh, sconcat("def p", poly_order, "c0(x, f, fL, fR, dx):~%"))$ -printf(fh, " return ~a", out) $ -close(fh)$ \ No newline at end of file diff --git a/src/postgkyl/data/select.py b/src/postgkyl/data/select.py deleted file mode 100644 index 848cff36..00000000 --- a/src/postgkyl/data/select.py +++ /dev/null @@ -1,97 +0,0 @@ - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -import postgkyl.data.idx_parser as idx_parser -if TYPE_CHECKING: - from postgkyl import GData -#end - - -def select(data: GData, comp: int | str | None = None, - z0: int | float | str | None = None, z1: int | float | str | None = None, - z2: int | float | str | None = None, z3: int | float | str | None = None, - z4: int | float | str | None = None, z5: int | float | str | None = None, - overwrite: bool = False) -> Tuple[list, np.ndarray]: - """Selects parts of the GData. - - Allows to select only a part of GData (both coordinates and - components). Allows for numpy slices, selecting multiple - components, and using both indicies (integer) and values (float). - - Atributes: - data (GData) - z0-5 (index, value, or slice (e.g. '1:5') - comp (index, slice (e.g. '1:5'), or multiple (e.g. '1,5') - """ - zs = (z0, z1, z2, z3, z4, z5) - grid = data.get_grid() - grid = list(grid) # copy the grid - values = data.get_values() - num_dims = data.get_num_dims() - bounds = data.get_bounds() - values_idx = [slice(0, values.shape[d]) for d in range(num_dims + 1)] - uniform_grid = len(grid[0].shape) == 1 - if not uniform_grid: - grid_idx = [slice(0, grid[d].shape[d]) for d in range(num_dims)] - # end - - # Loop for coordinates - for d, z in enumerate(zs): - if d < num_dims and z is not None: - #dat_range = bounds[1][d] - bounds[0][d] - #if '.' in z: - # if bounds[1][d] + 0.25 * dat_range < float(z) or bounds[0][d] - 0.25 * dat_range > float(z): - # raise TypeError("The coordinate select is outside of the data boundaries") - # #end - ##end - if uniform_grid: - len_grid = grid[d].shape[0] - else: - len_grid = grid[d].shape[d] - # end - is_matching = values.shape[d] == len_grid - idx = idx_parser.idx_parser(z, grid[d], is_matching) - if isinstance(idx, int): - # when 'slice' is used instead of an integer - # number, numpy array is not squeezed after - # subselecting - v_idx = slice(idx, idx + 1) - g_idx = slice(idx, idx + 2) if not is_matching else slice(idx, idx + 1) - elif isinstance(idx, slice): - v_idx = idx - g_idx = slice(idx.start, idx.stop + 1) if not is_matching else idx - else: - raise TypeError("The coordinate select can be only single index (int) or a slice.") - # end - if uniform_grid: - grid[d] = grid[d][g_idx] - else: - grid_idx[d] = g_idx - # end - values_idx[d] = v_idx - # end - # end - - # Select components - if comp is not None: - values_idx[-1] = idx_parser.idx_parser(comp) - # end - values_out = values[tuple(values_idx)] - if not uniform_grid: - for d in range(num_dims): - grid[d] = grid[d][tuple(grid_idx)] - # end - # end - - # Adding a dummy dimension indicies - if num_dims == len(values_out.shape): - values_out = values_out[..., np.newaxis] - # end - - if overwrite: - data.push(grid, values_out) - #end - return grid, values_out diff --git a/src/postgkyl/data/xformMatricesModalMaximal.h5 b/src/postgkyl/data/xformMatricesModalMaximal.h5 deleted file mode 100644 index 2d9cff4f..00000000 Binary files a/src/postgkyl/data/xformMatricesModalMaximal.h5 and /dev/null differ diff --git a/src/postgkyl/data/xformMatricesModalSerendipity.h5 b/src/postgkyl/data/xformMatricesModalSerendipity.h5 deleted file mode 100644 index f07a87e8..00000000 Binary files a/src/postgkyl/data/xformMatricesModalSerendipity.h5 and /dev/null differ diff --git a/src/postgkyl/data/xformMatricesNodalSerendipity.h5 b/src/postgkyl/data/xformMatricesNodalSerendipity.h5 deleted file mode 100644 index 23f8330d..00000000 Binary files a/src/postgkyl/data/xformMatricesNodalSerendipity.h5 and /dev/null differ diff --git a/src/postgkyl/dg/__init__.py b/src/postgkyl/dg/__init__.py new file mode 100644 index 00000000..65da201d --- /dev/null +++ b/src/postgkyl/dg/__init__.py @@ -0,0 +1,25 @@ +"""Discontinuous-Galerkin layer -- orchestrates Gkeyll's compiled DG engine. + +Four modules, one per domain boundary: + +- :mod:`.interpolate` -- the one-way modal -> NumPy bridge (matrix from + Gkeyll's basis functions, applied with NumPy); also ``local_poly``, the + same bridge evaluated at whole-cell points with NaN-separated interfaces + (the discontinuity-preserving plotting mesh). +- :mod:`.modal` -- operations that stay in the modal domain (weak algebra, + coefficient linear combinations, integration), all executed by Gkeyll + kernels on native arrays. +- :mod:`.rep` -- explicit value_form changes (modal · nodal · quad) and + pointwise functions via quadrature; the field never leaves the native domain. +- :mod:`.map` -- grid mapping: evaluate a coordinate-map field's coefficients + at a target's own grid points (see ``MAPPING.md``). +""" + +from .interpolate import interpolate, local_poly, num_basis +from .map import eval_at_points, map_grid, map_grid_separable +from . import modal, rep + +__all__ = [ + "interpolate", "local_poly", "num_basis", "modal", "rep", "eval_at_points", + "map_grid", "map_grid_separable" +] diff --git a/src/postgkyl/dg/interpolate.py b/src/postgkyl/dg/interpolate.py new file mode 100644 index 00000000..3eb0b5c8 --- /dev/null +++ b/src/postgkyl/dg/interpolate.py @@ -0,0 +1,181 @@ +"""Discontinuous-Galerkin interpolation -- modal coefficients -> mesh values. + +**This is the one-way bridge between the two domains**: DG coefficients in +(read through the container's NumPy view of the native array), plain NumPy +values out. The interpolation matrix is built from Gkeyll's own basis +functions (:mod:`postgkyl.gpython.basis` calls the ``eval`` pointer carried by +``struct gkyl_basis``), then applied per cell with a NumPy ``tensordot`` -- +so the result is always a *new, by-value* NumPy array, never a view of C +memory. The vendored sympy matrix tables this replaced lived in +``matrices.py`` (see ``src_bak`` history). + +:func:`local_poly` is the same bridge with a different evaluation-point +convention: points span the whole reference cell ``[-1, 1]`` (endpoints +included) instead of interior subcell centers, and a NaN is spliced in at +every cell interface -- so a plot shows the true DG inter-cell discontinuity +instead of the spuriously smooth curve :func:`interpolate` produces. The +hand-derived per-order polynomial tables the old implementation used +(``modalDG/kernels/expand_*d.py``, serendipity only) are superseded by +:func:`postgkyl.gpython.basis.eval_matrix`, which evaluates *any* basis at +arbitrary points through Gkeyll's own compiled basis-eval. +""" + +from __future__ import annotations + +import numpy as np + +from postgkyl.gpython import basis as gpython_basis + + +def num_basis(dim: int, poly_order: int, basis_type: str) -> int: + """Number of DG basis functions, straight from Gkeyll's basis object.""" + return gpython_basis.num_basis(basis_type, dim, poly_order) + + +def _make_mesh(num_interp: int, edges: np.ndarray) -> np.ndarray: + """Refine a 1-D nodal mesh by ``num_interp`` points per cell (uniform).""" + nx = edges.shape[0] - 1 + return np.linspace(edges[0], edges[-1], num_interp * nx + 1) + + +def _interpolate_on_mesh(c_mat: np.ndarray, q_in: np.ndarray, + num_interp: int) -> np.ndarray: + """Apply the interpolation matrix on every cell (per-point scatter).""" + num_cells = np.array(q_in.shape)[:-1] # drop the coefficient axis + num_dims = int(len(num_cells)) + ni = np.array([num_interp] * num_dims) + q_out = np.zeros(num_cells * ni, np.float64) + q_in = np.moveaxis(q_in, -1, 0) # coefficient index first + for n in range(int(np.prod(ni))): + temp = np.tensordot(c_mat[n, :], q_in, axes=1) + start_idx = np.unravel_index(n, ni, order="F") + idxs = [ + slice(int(start_idx[i]), int(num_cells[i] * ni[i]), int(ni[i])) + for i in range(num_dims) + ] + q_out[tuple(idxs)] = temp + return q_out + + +def interpolate(values: np.ndarray, + grid: list, + *, + poly_order: int, + basis_type: str, + nodal: bool = False, + num_interp: int | None = None): + """Interpolate DG coefficients onto a refined uniform mesh. + + Args: + values: ``(cells..., total_comps)`` array of DG coefficients. + grid: list of 1-D nodal edge arrays (one per dimension). + poly_order: polynomial order of the basis. + basis_type: long basis name (``"serendipity"``, ``"tensor"``, + ``"hybrid"``, or ``"gkhybrid"``). + nodal: Whether the data are field-blocked node values per cell; + converted through the exact ``nodal_to_modal`` matrix first. + num_interp: interpolation points per cell; defaults to ``poly_order + 1``. + + Returns: + ``(grid_out, values_out)`` -- the refined edge grid and a **new** + ``(refined_cells..., num_fields)`` NumPy value array. + """ + num_dims = len(grid) + if num_dims == 1 and basis_type == "hybrid": + basis_type = "serendipity" # PKPM hybrid degenerates to serendipity in 1D + if num_interp is None: + num_interp = poly_order + 1 + + nodes = num_basis(num_dims, poly_order, basis_type) + num_fields = values.shape[-1] // nodes + c_mat = gpython_basis.interpolation_matrix(basis_type, num_dims, poly_order, + num_interp) + + n2m = (gpython_basis.nodal_to_modal_matrix(basis_type, num_dims, poly_order) + if nodal else None) + out = None + for c in range(num_fields): + q = values[..., c * nodes:(c + 1) * nodes] + if n2m is not None: + q = np.einsum("jk,...k->...j", n2m, q) + interpolated_c = _interpolate_on_mesh(c_mat, q, num_interp)[..., np.newaxis] + out = interpolated_c if out is None else np.append( + out, interpolated_c, axis=-1) + + grid_out = [_make_mesh(num_interp, grid[d]) for d in range(num_dims)] + return grid_out, out + + +def _cell_edges_to_nodes(edges: np.ndarray, nodes_1d: np.ndarray) -> np.ndarray: + """Physical coordinates of ``nodes_1d`` (in ``[-1, 1]``) within every cell + of a 1-D ``edges`` array, flattened cell-major. Works for a non-uniform + grid: each cell is scaled/shifted from its own actual width, not assumed + uniform across the domain (unlike :func:`_make_mesh`).""" + cell_center = 0.5 * (edges[:-1] + edges[1:]) + dx = edges[1:] - edges[:-1] + return (cell_center[:, np.newaxis] + + nodes_1d[np.newaxis, :] * dx[:, np.newaxis] / 2).reshape(-1) + + +def local_poly(values: np.ndarray, + grid: list, + *, + poly_order: int, + basis_type: str, + nodal: bool = False, + npoints: int = 2): + """Evaluate the DG polynomial cell-by-cell onto a discontinuity-preserving + plotting mesh. + + Unlike :func:`interpolate`, points span the whole reference cell + ``[-1, 1]`` (``npoints`` of them, endpoints included) and a NaN is + inserted at every cell interface, so the true inter-cell jump of the DG + solution is visible when plotted instead of hidden by a spuriously smooth + curve. + + Args: + values: ``(cells..., total_comps)`` array of DG coefficients. + grid: list of 1-D nodal edge arrays (one per dimension). + poly_order: polynomial order of the basis. + basis_type: long basis name (``"serendipity"``, ``"tensor"``, + ``"hybrid"``, or ``"gkhybrid"``). + nodal: Whether the data use a nodal basis; converted through the exact + ``nodal_to_modal`` matrix first. + npoints: evaluation points per cell, from one face to the other. + + Returns: + ``(grid_out, values_out)`` -- a NaN-separated edge-grid list and value + array, one entry longer per cell interface than the plain ``npoints`` + x ``num_cells`` mesh. + """ + num_dims = len(grid) + if num_dims == 1 and basis_type == "hybrid": + basis_type = "serendipity" # PKPM hybrid degenerates to serendipity in 1D + + nodes_1d = np.linspace(-1.0, 1.0, npoints) + num_nodes = len(nodes_1d) + + nb = num_basis(num_dims, poly_order, basis_type) + num_fields = values.shape[-1] // nb + c_mat = gpython_basis.eval_matrix( + basis_type, num_dims, poly_order, + gpython_basis.tensor_points(nodes_1d, num_dims)) + + n2m = (gpython_basis.nodal_to_modal_matrix(basis_type, num_dims, poly_order) + if nodal else None) + out = None + for c in range(num_fields): + q = values[..., c * nb:(c + 1) * nb] + if n2m is not None: + q = np.einsum("jk,...k->...j", n2m, q) + field_c = _interpolate_on_mesh(c_mat, q, num_nodes)[..., np.newaxis] + out = field_c if out is None else np.append(out, field_c, axis=-1) + + num_cells = np.array(values.shape[:-1]) + grid_out = [_cell_edges_to_nodes(grid[d], nodes_1d) for d in range(num_dims)] + for d in range(num_dims): + sep = np.arange(num_nodes, num_nodes * num_cells[d], num_nodes) + out = np.insert(out, sep, np.nan, axis=d) + grid_out[d] = np.insert(grid_out[d], sep, grid_out[d][sep - 1]) + + return grid_out, out diff --git a/src/postgkyl/dg/map.py b/src/postgkyl/dg/map.py new file mode 100644 index 00000000..49127dfd --- /dev/null +++ b/src/postgkyl/dg/map.py @@ -0,0 +1,199 @@ +"""Grid mapping -- evaluate a coordinate-map DG field at a target's grid points. + +See ``MAPPING.md`` for the full design. **The core semantic**: a mapping file +is a DG field whose components hold the coefficients of the physical +coordinates :math:`x_d(z)` of each mapped dimension ``d``; mapping a grid means +evaluating those coefficients at the *target*'s existing grid points -- there is +no resolution parameter and no alignment arithmetic, so the new grid always has +exactly the shape of the one it replaces. + +Two functions, one per step of ``MAPPING.md``'s "evaluation algorithm": +:func:`eval_at_points` evaluates ONE coordinate's coefficients at an arbitrary +point set (steps 1-4: cell locate, reference-coordinate conversion, grouped +basis evaluation, reshape); :func:`map_grid` builds the tensor point set for +the target axes and calls it once per mapped dimension. +""" + +from __future__ import annotations + +import numpy as np + +from postgkyl import gpython + + +def eval_at_points(coeffs: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + cells: np.ndarray, + points: np.ndarray, + *, + basis_type: str, + poly_order: int, + nodal: bool = False) -> np.ndarray: + """Evaluate one coordinate's DG coefficients at arbitrary computational points. + + Args: + coeffs: ``(*cells, num_basis)`` array -- the mapping's per-cell + coefficients for a single physical coordinate ``x_d(z)`` over its own + uniform grid. Pass ``nodal=True`` for a nodal-basis mapping file + (converted through the exact ``nodal_to_modal`` matrix first, same + pattern as :func:`postgkyl.dg.interpolate.interpolate`). + lower: length-``m`` array, the mapping's own domain lower bounds. + upper: length-``m`` array, the mapping's own domain upper bounds. + cells: length-``m`` array, the mapping's own cell counts (must match + ``coeffs.shape[:-1]``). + points: ``(*shape, m)`` array of evaluation points in computational + coordinates, within the mapping's bounds. + basis_type: long basis name (``"serendipity"``, ``"tensor"``, + ``"hybrid"``, or ``"gkhybrid"``). + poly_order: polynomial order of the mapping's basis. + nodal: Whether these are nodal-basis mapping coefficients. + + Returns: + ``(*shape,)`` array -- ``x_d`` evaluated at every point. + + Raises: + ValueError: ``cells`` does not match ``coeffs.shape[:-1]``, or the last + axis of ``points`` does not have length ``m``. + """ + lower = np.asarray(lower, dtype=np.float64) + upper = np.asarray(upper, dtype=np.float64) + cells = np.asarray(cells, dtype=np.int64) + m = lower.shape[0] + + if coeffs.shape[:-1] != tuple(int(c) for c in cells): + raise ValueError( + f"eval_at_points: coeffs cell shape {coeffs.shape[:-1]} does not " + f"match cells {tuple(cells)}") + points = np.asarray(points, dtype=np.float64) + if points.shape[-1] != m: + raise ValueError( + f"eval_at_points: points last axis has length {points.shape[-1]}, " + f"expected {m} (len(lower))") + + if nodal: + n2m = gpython.basis.nodal_to_modal_matrix(basis_type, m, poly_order) + coeffs = np.einsum("jk,...k->...j", n2m, coeffs) + + shape = points.shape[:-1] + z = points.reshape(-1, m) + dz = (upper - lower) / cells + + # Step 1: locate the containing cell (clip fixes the boundary convention). + idx = np.clip(np.floor((z - lower) / dz).astype(np.int64), 0, cells - 1) + # Step 2: reference coordinates in [-1, 1]^m. + centers = lower + (idx + 0.5) * dz + eta = 2.0 * (z - centers) / dz + + flat_coeffs = coeffs.reshape(-1, coeffs.shape[-1]) + cell_lin = np.ravel_multi_index(tuple(idx[:, d] for d in range(m)), + tuple(int(c) for c in cells)) + + # Step 3: group points by containing cell -> one matrix-vector product each. + out = np.empty(z.shape[0], dtype=np.float64) + for lin in np.unique(cell_lin): + sel = cell_lin == lin + b = gpython.basis.eval_matrix(basis_type, m, poly_order, eta[sel]) + out[sel] = b @ flat_coeffs[lin] + + # Step 4: reshape to the point-set shape. + return out.reshape(shape) + + +def map_grid(map_coeffs: np.ndarray, map_ctx: dict, + target_axes: list[np.ndarray]) -> list[np.ndarray]: + """Evaluate every mapped dimension's coordinates at the target's grid points. + + Args: + map_coeffs: ``(*cells, m * num_basis)`` array -- the mapping field's raw + coefficients (``GDataState.get_values()``); components + ``d*num_basis:(d+1)*num_basis`` are the coefficients of ``x_d(z)``. + map_ctx: the mapping dataset's ``ctx`` dict; reads ``lower``, ``upper``, + ``cells``, ``basis_type``, ``poly_order``, and ``value_form`` (default + ``"modal"``). + target_axes: the target's own edge/node arrays for the ``m`` axes being + deformed, one 1-D array per mapped dimension. + + Returns: + A list of ``m`` new grid arrays, one per mapped dimension: 1-D when + ``m == 1``; an ``m``-dimensional nodal array (the full tensor product of + ``target_axes``, ``indexing="ij"``) for every dimension when ``m > 1``, + so non-separable (curvilinear) maps are handled. + """ + lower = map_ctx["lower"] + upper = map_ctx["upper"] + cells = map_ctx["cells"] + basis_type = map_ctx["basis_type"] + poly_order = map_ctx["poly_order"] + nodal = map_ctx.get("value_form", "modal") == "nodal" + m = len(target_axes) + + if m == 1: + points = np.asarray(target_axes[0], dtype=np.float64)[:, np.newaxis] + else: + points = np.stack(np.meshgrid(*target_axes, indexing="ij"), axis=-1) + + nb = gpython.basis.num_basis(basis_type, m, poly_order) + return [ + eval_at_points(map_coeffs[..., d * nb:(d + 1) * nb], + lower, + upper, + cells, + points, + basis_type=basis_type, + poly_order=poly_order, + nodal=nodal) for d in range(m) + ] + + +def map_grid_separable(map_coeffs: np.ndarray, map_ctx: dict, + target_axes: list[np.ndarray]) -> list[np.ndarray]: + """Evaluate each mapped dimension's coordinates independently. + + Gkeyll's velocity-space coordinate maps (``mapc2p_vel``) are diagonal: + dimension ``d`` is a **separate** 1-D map ``v_d(z_d)`` over its own axis + only, not a joint ``m``-dimensional curvilinear map like :func:`map_grid` + handles for configuration-space maps (``mapc2p``/``mc2nu``). Component + block ``d*num_basis:(d+1)*num_basis`` (``num_basis`` for a *1-D* basis of + the given order) holds dimension ``d``'s own coefficients; Gkeyll's writer + stores them on the full ``m``-dimensional cell grid but broadcasts them + along every axis other than ``d``, so cell index 0 on every other axis + already carries the full set of values. + + Args: + map_coeffs: ``(*cells, m * num_basis)`` array -- ``GDataState.get_values()``, + ``num_basis`` being the 1-D basis size for ``basis_type``/``poly_order``. + map_ctx: the mapping dataset's ``ctx`` dict; reads ``lower``, ``upper``, + ``cells``, ``basis_type``, ``poly_order``, and ``value_form`` (default + ``"modal"``). + target_axes: the target's own 1-D edge arrays for the ``m`` axes being + deformed, one per mapped dimension. + + Returns: + A list of ``m`` new 1-D grid arrays, one per mapped dimension. + """ + lower = map_ctx["lower"] + upper = map_ctx["upper"] + cells = map_ctx["cells"] + basis_type = map_ctx["basis_type"] + poly_order = map_ctx["poly_order"] + nodal = map_ctx.get("value_form", "modal") == "nodal" + m = len(target_axes) + + nb = gpython.basis.num_basis(basis_type, 1, poly_order) + new_axes = [] + for d in range(m): + idx = [0] * m + idx[d] = slice(None) + coeffs_d = map_coeffs[tuple(idx)][:, d * nb:(d + 1) * nb] + points = np.asarray(target_axes[d], dtype=np.float64)[:, np.newaxis] + new_axes.append( + eval_at_points(coeffs_d, + lower[d:d + 1], + upper[d:d + 1], + cells[d:d + 1], + points, + basis_type=basis_type, + poly_order=poly_order, + nodal=nodal)) + return new_axes diff --git a/src/postgkyl/dg/modal.py b/src/postgkyl/dg/modal.py new file mode 100644 index 00000000..26cf153f --- /dev/null +++ b/src/postgkyl/dg/modal.py @@ -0,0 +1,259 @@ +"""Modal (DG-coefficient) operations -- thin orchestration over Gkeyll kernels. + +Everything here acts on native :class:`~postgkyl.gpython.array.GkylArray` data and +returns native data (or plain numbers for reductions): the modal domain never +leaves Gkeyll's memory. The only logic this layer adds over ``gpython.kernels`` is +DG bookkeeping -- e.g. what "add a scalar" means for modal coefficients. +""" + +from __future__ import annotations + +import numpy as np + +from postgkyl import gpython +from postgkyl.gpython.array import GkylArray + +# Weak algebra and coefficient linear combinations -- direct kernel calls. +weak_mul = gpython.kernels.weak_mul +weak_div = gpython.kernels.weak_div +weak_inv = gpython.kernels.weak_inv +weak_mul_conf_phase = gpython.kernels.weak_mul_conf_phase +lincomb = gpython.kernels.lincomb +scale = gpython.kernels.scale +integrate = gpython.kernels.integrate +reduce = gpython.kernels.reduce + + +def is_native(value) -> bool: + """True if ``value`` is a native (gkyl-backed) array, not a plain NumPy one. + + The one place outside ``gdatastate``/``gpython`` that needs to tell modal + data apart from plain arrays without importing ``gpython`` directly (an + import-contract boundary; see ``operations.evaluate``). + """ + return isinstance(value, GkylArray) + + +def shift_mean(basis_type: str, ndim: int, poly_order: int, a: GkylArray, + val: float) -> GkylArray: + """``f + val`` for a modal field: only the mean coefficient moves. + + The normalized constant basis function is ``b_0 = 2^(-ndim/2)``, so a shift + of the field by ``val`` is a shift of coefficient 0 by ``val * 2^(ndim/2)``, + applied per field (``gkyl_array_shiftc`` on each field's coefficient 0). + """ + nb = gpython.basis.num_basis(basis_type, ndim, poly_order) + coeff_shift = float(val) * 2.0**(ndim / 2.0) + out = a + for f in range(a.ncomp // nb): + out = gpython.kernels.shiftc(out, coeff_shift, f * nb) + return out + + +def shift_all(a: GkylArray, val: float) -> GkylArray: + """``values + val`` for point-value representations (nodal/quad): every + component of every cell is a field value, so shift them all.""" + out = a.clone() + for k in range(a.ncomp): + out = gpython.kernels.shiftc(out, float(val), k) + return out + + +def average(grid: dict, + basis_type: str, + ndim: int, + poly_order: int, + a: GkylArray, + avg_dirs, + weight: GkylArray | None = None): + """``int f w dx^avg / int w dx^avg`` (or the plain average) of a modal + field over ``avg_dirs``, field by field. + + ``gkyl_array_average`` has no field-index argument (unlike the weak ops + above, which loop inside the compiled shim), so a multi-field ``a`` + (``ncomp == nfields * num_basis``) is split into single-field slices here + and averaged one at a time, then reassembled. + + Args: + grid: donor grid dict (``ndim``/``lower``/``upper``/``cells``, e.g. from + ``rio``). + avg_dirs: 0-based donor directions to average over. + weight: optional single-field ``GkylArray`` over the same donor + grid/basis as ``a`` (the plain average, dividing by volume, is used + when omitted). + + Returns: + ``(keep_dirs, cells_avg, result)`` -- the surviving donor directions (in + order), the target's per-dimension cell counts, and the averaged array + (``ncomp`` scaled to the same field count as ``a``). ``keep_dirs``/ + ``cells_avg`` are empty/``[1]`` for a full reduction: Gkeyll always + keeps at least one target dimension, collapsing to a single cell when + every donor direction is averaged out. + """ + avg_dirs = sorted(set(int(d) for d in avg_dirs)) + if not avg_dirs or avg_dirs[0] < 0 or avg_dirs[-1] >= ndim: + raise ValueError( + f"average dirs {avg_dirs} out of range for a {ndim}D field") + keep_dirs = [d for d in range(ndim) if d not in avg_dirs] + ndim_avg = len(keep_dirs) if keep_dirs else 1 + cells = np.asarray(grid["cells"]) + cells_avg = [int(cells[d]) for d in keep_dirs] if keep_dirs else [1] + avg_dim = [1 if d in avg_dirs else 0 for d in range(ndim)] + + nb = gpython.basis.num_basis(basis_type, ndim, poly_order) + if a.ncomp % nb: + raise ValueError(f"ncomp {a.ncomp} is not a multiple of num_basis {nb}") + nfields = a.ncomp // nb + if weight is not None and weight.ncomp != nb: + raise ValueError(f"average weight ncomp ({weight.ncomp}) must equal " + f"the donor basis's num_basis ({nb})") + + if nfields == 1: + out = gpython.kernels.array_average(grid, + basis_type, + poly_order, + ndim_avg, + cells_avg, + avg_dim, + a, + weight=weight) + else: + a_view = a.view().reshape(a.size, nfields, nb) + fields_out = [] + for f in range(nfields): + a_f = GkylArray.from_numpy(np.ascontiguousarray(a_view[:, f, :])) + out_f = gpython.kernels.array_average(grid, + basis_type, + poly_order, + ndim_avg, + cells_avg, + avg_dim, + a_f, + weight=weight) + fields_out.append(out_f.view()) + out = GkylArray.from_numpy(np.concatenate(fields_out, axis=-1)) + + if not keep_dirs and weight is None: + # Full reduction (every donor dim averaged), unweighted only: Gkeyll's + # own kernels for this corner case (gkyl_array_average_NxYY_avg) write a single raw VALUE into coefficient 0 -- there is no real + # target dimension to normalize against, unlike every other path here + # (a partial reduction, or ANY weighted reduction, which both go through + # a genuine per-mode contraction/weak-division and so already come out + # as a properly b0-normalized coefficient). Rescale so this dataset's + # coefficient 0 means the same thing ("value = coeff0 * b0") as every + # other modal dataset in the system -- verified against + # gkyl_array_integrate on a constant field (see test_dg_modal_average). + out = gpython.kernels.scale(out, 2.0**(ndim_avg / 2.0)) + return keep_dirs, cells_avg, out + + +def differentiate(basis_type: str, ndim: int, poly_order: int, a: GkylArray, + dir: int, diff_order: int, dx: float) -> GkylArray: + """``d^diff_order/dx_dir^diff_order a``, field by field + (``gkyl_dg_differentiate_op_local`` -- exact on the polynomial each cell + already represents; no inter-cell stencil). The field loop lives in the + shim (like :func:`weak_mul`), so this is a direct pass-through.""" + return gpython.kernels.weak_differentiate(basis_type, ndim, poly_order, dir, + diff_order, dx, a) + + +def eval_at_coord_proj(grid: dict, basis_type: str, ndim: int, poly_order: int, + a: GkylArray, eval_dirs, eval_coords): + """Evaluate a modal field at ``eval_coords`` in ``eval_dirs`` and project + onto the surviving directions' target basis (``gkyl_dg_eval_at_coord_proj``). + + ``grid`` is the donor grid dict (``ndim``/``lower``/``upper``/``cells``, + e.g. from ``rio``). The donor's configuration-space dimension count (needed + by the underlying updater) is derived from ``basis_type``/``ndim`` via + ``gpython.basis.cdim_vdim``. + + Returns: + ``(keep_dirs, cells_tar, out, target_basis_type, target_poly_order, + target_cdim, target_vdim)`` -- ``keep_dirs``/``cells_tar`` follow the + same full-reduction convention :func:`average` uses (empty/``[1]`` when + every donor direction is evaluated away, since Gkeyll always keeps at + least one target dimension). + """ + eval_dirs = sorted(set(int(d) for d in eval_dirs)) + if not eval_dirs or eval_dirs[0] < 0 or eval_dirs[-1] >= ndim: + raise ValueError(f"eval_dirs {eval_dirs} out of range for a {ndim}D field") + keep_dirs = [d for d in range(ndim) if d not in eval_dirs] + cells = np.asarray(grid["cells"]) + ndim_tar = len(keep_dirs) if keep_dirs else 1 + cells_tar = [int(cells[d]) for d in keep_dirs] if keep_dirs else [1] + cdim_do, _vdim_do = gpython.basis.cdim_vdim(basis_type, ndim) + + out, btype, poly_order_tar, cdim_tar, vdim_tar = ( + gpython.kernels.eval_at_coord_proj(basis_type, ndim, poly_order, cdim_do, + grid, eval_dirs, eval_coords, ndim_tar, + cells_tar, a)) + return (keep_dirs, cells_tar, out, btype, poly_order_tar, cdim_tar, vdim_tar) + + +def power(basis_type: str, + ndim: int, + poly_order: int, + a: GkylArray, + exponent, + cells=None) -> GkylArray: + """``f ** n``. + + A positive integer ``n`` takes the cheap, exact path: repeated weak + multiplies. Any other exponent (0, negative, or fractional) falls + through to :func:`powsqrt` (``f ** n == pow(sqrt(f), 2n)``), which needs + ``cells`` (the grid's per-dimension cell count, e.g. ``ctx["cells"]``) to + build Gkeyll's index range -- required whenever this fallback fires. + """ + n = exponent + if isinstance(n, (int, np.integer)) and n >= 1: + out = a.clone() + for _ in range(int(n) - 1): + out = weak_mul(basis_type, ndim, poly_order, out, a) + return out + if cells is None: + raise ValueError( + f"modal power with exponent {n!r} (not a positive integer) needs " + "cells= to build the powsqrt kernel's index range.") + return powsqrt(basis_type, ndim, poly_order, cells, a, 2.0 * float(n)) + + +def powsqrt(basis_type: str, + ndim: int, + poly_order: int, + cells, + a: GkylArray, + exponent: float, + num_quad: int | None = None) -> GkylArray: + """``pow(sqrt(f), exponent)`` (i.e. ``f ** (exponent/2)``), field by field. + + ``gkyl_proj_powsqrt_on_basis`` has no field-index argument (like + :func:`average`'s ``gkyl_array_average``), so a multi-field ``a`` + (``ncomp == nfields * num_basis``) is split into single-field slices + here and processed one at a time, then reassembled. + """ + nb = gpython.basis.num_basis(basis_type, ndim, poly_order) + if a.ncomp % nb: + raise ValueError(f"ncomp {a.ncomp} is not a multiple of num_basis {nb}") + nfields = a.ncomp // nb + if nfields == 1: + return gpython.kernels.powsqrt(basis_type, + ndim, + poly_order, + cells, + a, + exponent, + num_quad=num_quad) + a_view = a.view().reshape(a.size, nfields, nb) + fields_out = [] + for f in range(nfields): + a_f = GkylArray.from_numpy(np.ascontiguousarray(a_view[:, f, :])) + out_f = gpython.kernels.powsqrt(basis_type, + ndim, + poly_order, + cells, + a_f, + exponent, + num_quad=num_quad) + fields_out.append(out_f.view()) + return GkylArray.from_numpy(np.concatenate(fields_out, axis=-1)) diff --git a/src/postgkyl/dg/rep.py b/src/postgkyl/dg/rep.py new file mode 100644 index 00000000..8fb390e4 --- /dev/null +++ b/src/postgkyl/dg/rep.py @@ -0,0 +1,193 @@ +"""Representation changes within the native domain -- modal · nodal · quad. + +One DG field, three per-cell representations (REFACTOR_GKEYLL_FFI.md §3b): +modal coefficients, values at the basis nodes, values at Gauss–Legendre +quadrature points. Conversions are per-cell matrix applications built from +Gkeyll's basis function pointers (:mod:`postgkyl.gpython.basis`); data enters and +leaves as a native :class:`~postgkyl.gpython.array.GkylArray`, so the field never +leaves the native domain. **Nothing here converts implicitly** -- these are the +backends of the explicit ``.to_nodal()/.to_modal()/.to_quad()/.apply()`` verbs. + +Exactness: nodal↔modal is an exact N×N change of basis; a quad round-trip is +exact for integrands of degree ≤ 2·num_quad−1 (default ``num_quad = p+1``). + +Note: this is the *cell-local* nodal value_form (N unshared values per +cell). Grid-level shared-node nodal fields (``gkyl_nodal_ops``, used by the +geometry/mapped-grid workflow) are phase C. +""" + +from __future__ import annotations + +import numpy as np + +from postgkyl.gpython import basis as gpython_basis +from postgkyl.gpython.array import GkylArray + + +def _apply_per_field(arr: GkylArray, comps_in: int, + mat: np.ndarray) -> GkylArray: + """Apply ``mat`` (comps_out × comps_in) to every field of every cell.""" + if arr.ncomp % comps_in: + raise ValueError(f"ncomp {arr.ncomp} is not a multiple of {comps_in}") + nfields = arr.ncomp // comps_in + v = arr.view().reshape(arr.size, nfields, comps_in) + out = np.einsum("pk,cfk->cfp", mat, v).reshape(arr.size, + nfields * mat.shape[0]) + return GkylArray.from_numpy(out) + + +def modal_to_nodal(basis_type: str, ndim: int, poly_order: int, + arr: GkylArray) -> GkylArray: + """Coefficients -> values at the basis ``node_list`` points (exact).""" + nb = gpython_basis.num_basis(basis_type, ndim, poly_order) + return _apply_per_field( + arr, nb, gpython_basis.modal_to_nodal_matrix(basis_type, ndim, + poly_order)) + + +def nodal_to_modal(basis_type: str, ndim: int, poly_order: int, + arr: GkylArray) -> GkylArray: + """Values at the basis nodes -> coefficients (exact inverse).""" + nb = gpython_basis.num_basis(basis_type, ndim, poly_order) + return _apply_per_field( + arr, nb, gpython_basis.nodal_to_modal_matrix(basis_type, ndim, + poly_order)) + + +def modal_to_quad(basis_type: str, ndim: int, poly_order: int, arr: GkylArray, + num_quad: int) -> GkylArray: + """Coefficients -> values at the tensor Gauss–Legendre points.""" + nb = gpython_basis.num_basis(basis_type, ndim, poly_order) + return _apply_per_field( + arr, nb, + gpython_basis.modal_to_quad_matrix(basis_type, ndim, poly_order, + num_quad)) + + +def quad_to_modal(basis_type: str, ndim: int, poly_order: int, arr: GkylArray, + num_quad: int) -> GkylArray: + """Quadrature values -> coefficients (projection; exact for degree + ≤ 2·num_quad−1).""" + nq = num_quad**ndim + return _apply_per_field( + arr, nq, + gpython_basis.quad_to_modal_matrix(basis_type, ndim, poly_order, + num_quad)) + + +def wrap(values: np.ndarray) -> GkylArray: + """Wrap ``(cells..., ncomp)`` NumPy values back into a native array. + + The doorway for pointwise NumPy results on nodal/quad data: computed on the + view, wrapped back, so the dataset stays gkyl-native and in-value_form. + """ + return GkylArray.from_numpy(values) + + +def _tensor_point_layout(basis_type: str, ndim: int, poly_order: int, rep: str, + num_quad: int | None): + """Per-dimension reference points + permutation into Fortran tensor order. + + Returns ``(pts_1d_per_dim, perm)`` where ``values[..., perm]`` reorders a + cell's point values into F-order tensor indexing (dimension 0 fastest). + Quadrature points are a tensor product by construction; nodal sets are + checked -- non-tensor node sets (e.g. serendipity p2 in 2-D+) raise. + """ + if rep == "quad": + nq = int(num_quad) if num_quad else poly_order + 1 + pts_1d, _ = np.polynomial.legendre.leggauss(nq) + return [pts_1d] * ndim, None + coords = gpython_basis.node_coords(basis_type, ndim, poly_order) + nb = coords.shape[0] + uniq = [np.unique(coords[:, d]) for d in range(ndim)] + counts = [len(u) for u in uniq] + if int(np.prod(counts)) != nb: + raise ValueError( + f"the {basis_type} p{poly_order} {ndim}D node set is not a tensor " + "product; use .to_quad() for point-value work in this basis.") + lin = np.zeros(nb, dtype=np.int64) + stride = 1 + for d in range(ndim): + k = np.searchsorted(uniq[d], coords[:, d]) + if not np.allclose(uniq[d][k], coords[:, d]): + raise ValueError("node coordinates do not align on a tensor grid") + lin += k * stride + stride *= counts[d] + if len(np.unique(lin)) != nb: + raise ValueError( + f"the {basis_type} p{poly_order} {ndim}D node set is not a tensor " + "product; use .to_quad() for point-value work in this basis.") + return [uniq[d] for d in range(ndim)], np.argsort(lin) + + +def _edges_from_points(pts: np.ndarray, lo: float, hi: float) -> np.ndarray: + """Edges such that cell centers coincide with ``pts`` (honest positions).""" + e = np.empty(len(pts) + 1) + e[0] = lo + for i in range(len(pts)): + e[i + 1] = 2.0 * pts[i] - e[i] + e[-1] = hi + return np.maximum.accumulate(e) # degenerate (zero-width) cells allowed + + +def materialize(basis_type: str, + ndim: int, + poly_order: int, + arr: GkylArray, + grid: list, + rep: str, + num_quad: int | None = None): + """Point-value data -> ``(nonuniform edge grid, ndarray)`` at the TRUE + physical point locations -- the render path for nodal/quad datasets. + + Unlike ``interpolate`` (which evaluates modal data on an equispaced mesh), + this performs no basis math at all: the values *are* the field at their + points; only coordinates and ordering are computed. + """ + pts_1d, perm = _tensor_point_layout(basis_type, ndim, poly_order, rep, + num_quad) + counts = [len(p) for p in pts_1d] + npc = int(np.prod(counts)) + if arr.ncomp % npc: + raise ValueError( + f"ncomp {arr.ncomp} is not a multiple of {npc} points/cell") + nfields = arr.ncomp // npc + cells = [len(g) - 1 for g in grid] + v = arr.view().reshape(*cells, nfields, npc) + if perm is not None: + v = v[..., perm] + + out = np.zeros([cells[d] * counts[d] for d in range(ndim)] + [nfields]) + for n in range(npc): + off = np.unravel_index(n, counts, order="F") + idxs = tuple( + slice(int(off[d]), cells[d] * counts[d], counts[d]) + for d in range(ndim)) + out[idxs] = v[..., n] + + edges = [] + for d in range(ndim): + g = np.asarray(grid[d], dtype=np.float64) + centers, dxs = 0.5 * (g[:-1] + g[1:]), np.diff(g) + pts = (centers[:, None] + 0.5 * dxs[:, None] * pts_1d[d][None, :]).ravel() + edges.append(_edges_from_points(pts, g[0], g[-1])) + return edges, out + + +def apply_pointwise(basis_type: str, ndim: int, poly_order: int, arr: GkylArray, + fn, num_quad: int) -> GkylArray: + """``fn`` applied pointwise via quadrature: modal → quad → fn → modal. + + The standard DG treatment of nonlinear operations. ``fn`` receives the + ``(cells, nfields*nq)`` array of quadrature values and must return the same + shape (any NumPy ufunc qualifies). The result is modal again. + """ + quad = modal_to_quad(basis_type, ndim, poly_order, arr, num_quad) + vals = fn(quad.view()) + vals = np.asarray(vals, dtype=np.float64) + if vals.shape != (quad.size, quad.ncomp): + raise ValueError( + f"apply(fn): fn changed the shape {(quad.size, quad.ncomp)} -> " + f"{vals.shape}; it must act pointwise.") + return quad_to_modal(basis_type, ndim, poly_order, GkylArray.from_numpy(vals), + num_quad) diff --git a/src/postgkyl/diagnostics/__init__.py b/src/postgkyl/diagnostics/__init__.py new file mode 100644 index 00000000..21a1a857 --- /dev/null +++ b/src/postgkyl/diagnostics/__init__.py @@ -0,0 +1,170 @@ +"""Equation-specific physics grouped by Gkeyll model family. + +Folds together the old ``models`` (array math) and ``operations`` physics-verb +(GData wrapping) layers into a single home per equation system: functions +here take loaded ``GData``/``GDataState`` (one or several) plus physical +scalars as keyword-only options, and return a ``GDataState`` (via +``_result``) or, in later layers, a ``Figure``. Equation-blind core verbs +stay in flat ``operations`` modules. Domain-specific transformations live in +operation subpackages (for example ``operations.gyrokinetics``); this layer +is reserved for code that knows what field components physically mean. + +The four public packages mirror Gkeyll's model families: ``gk``, ``vm``, +``pkpm``, and ``mom``. The equation-blind ``discovery`` module +stays shared at this package root. There is no separate ``loaders`` package; +each model family owns its loading and program-scale diagnostics. +""" + +from . import discovery, gk, mom, pkpm, vm + +from typing import Annotated, Literal + +from postgkyl.cli_spec import ( + CommandSpec, + DatasetRef, + Execution, + ResultPolicy, + Section, + command, + command_spec, + hidden, + hidden_spec, +) +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.gdata.gdata import GData + +_DIAG_MAP = CommandSpec(Section.DIAGNOSTICS, Execution.MAP_REPLACE) +_DIAG_COMBINE = CommandSpec(Section.DIAGNOSTICS, + Execution.COMBINE, + consumes_inputs=True) +_DIAG_LOAD = CommandSpec(Section.DIAGNOSTICS, Execution.LOAD) +_DIAG_REPORT = CommandSpec(Section.DIAGNOSTICS, + Execution.LOAD, + result=ResultPolicy.VALUE) + + +def _resolve(function) -> None: + function.__globals__.setdefault("GDataState", GDataState) + function.__globals__.setdefault("_GDataState", GDataState) + function.__globals__.setdefault("GData", GData) + + +def _map(function) -> None: + _resolve(function) + command(_DIAG_MAP)(function) + + +def _combine(function, dataset_names: tuple[str, ...]) -> None: + _resolve(function) + for name in dataset_names: + function.__annotations__[name] = Annotated[GDataState, DatasetRef()] + command(_DIAG_COMBINE)(function) + + +for _module, _names in ( + (mom.five_moment, ("density", "xvel", "yvel", "zvel", "vel", "pressure", + "ke", "temp", "sound", "mach")), + (mom.ten_moment, ("pressure", "ke", "temp", "sound", "mach", "pxx", "pxy", + "pxz", "pyy", "pyz", "pzz", "pressure_tensor")), + (mom.mhd, ("bx", "by", "bz", "bi", "mag_pressure", "pressure", "temp", + "sound", "mach")), + (mom.plasma, ("magB", "vt", "omegaC", "omegaP", "d", "lambdaD")), +): + for _name in _names: + _map(getattr(_module, _name)) + +_resolve(mom.multispecies.accumulate_current) +command( + CommandSpec(Section.DIAGNOSTICS, Execution.MAP_APPEND, + consumes_inputs=True))(mom.multispecies.accumulate_current) + +for _function, _datasets in ( + (mom.five_moment.velocity, ("density", "momentum")), + (mom.ten_moment.p_par, ("ptensor", "bfield")), + (mom.ten_moment.p_perp, ("ptensor", "bfield")), + (mom.ten_moment.agyro, ("ptensor", "bfield")), + (mom.ten_moment.mom_agyro, ("species", "field")), + (mom.plasma.vA, ("species", "field")), + (mom.plasma.rho, ("species", "field")), + (mom.plasma.beta, ("species", "field")), + (mom.multispecies.energetics, ("elc", "ion", "field")), + (mom.rotations.parrotate, ("array", "rotator")), + (mom.rotations.perprotate, ("array", "rotator")), + (mom.rotations.bparrotate, ("array", "field")), + (mom.rotations.bperprotate, ("array", "field")), + (vm.kinetic.transform_frame, ("distribution", "bulk")), + (pkpm.laguerre_compose, ("distribution", "variables")), +): + _combine(_function, _datasets) + +mom.ten_moment.agyro.__annotations__["measure"] = Literal["swisdak", + "frobenius"] +mom.ten_moment.mom_agyro.__annotations__["measure"] = Literal["swisdak", + "frobenius"] + +for _function in (pkpm.load_pkpm, discovery.find_output_stems, + discovery.available_frames, mom.enstrophy.enstrophy, + mom.ke_dke.ke_dke): + _resolve(_function) + command(_DIAG_LOAD if _function is pkpm.load_pkpm else _DIAG_REPORT)( + _function) +pkpm.load_pkpm.__annotations__["idx"] = str + +_resolve(vm.trajectory.trajectory) +command( + CommandSpec(Section.DIAGNOSTICS, + Execution.TERMINAL_ALL, + result=ResultPolicy.VALUE))(vm.trajectory.trajectory) + +for _function in ( + gk.load_distf, + gk.load_quantity, + gk.energy_balance, + gk.particle_balance, + gk.nodes, +): + _resolve(_function) + +# These exclusions are an explicit audit, not a catch-all over module +# contents. Adding a public diagnostic callable without classifying it now +# fails discovery instead of being silently hidden. +for _name in ( + "resolve_frames", + "available_quantities", + "fetch_beta_from_bmag_press", + "fetch_diamag_vel", + "fetch_ExB_vel", + "fetch_gradB_vel", + "fetch_M1_from_H", + "fetch_press_from_BiMax", + "fetch_press_from_Max", + "fetch_press_p", + "fetch_Tpar_from_BiMax", + "fetch_Tpar_from_M0_M1_M2par", + "fetch_temp_from_Max", + "fetch_temp_from_Tpar_Tperp", + "fetch_Tperp_from_BiMax", + "fetch_Tperp_from_M0_M2perp", + "energy_balance_error", + "particle_balance_error", + "is_geo_mapc2p", + "multib_tag", + "nodes_to_RZ", + "map_to_rz", + "resolve_geometry", + "resolve_rz_projection", + "extract_flux_surface", + "resolve_flux_surface_grid", +): + _function = getattr(gk, _name) + if command_spec(_function) is None and hidden_spec(_function) is None: + hidden("requires Python objects or is a registry/provider helper")( + _function) + +__all__ = [ + "gk", + "vm", + "mom", + "pkpm", + "discovery", +] diff --git a/src/postgkyl/diagnostics/discovery.py b/src/postgkyl/diagnostics/discovery.py new file mode 100644 index 00000000..3aead7c2 --- /dev/null +++ b/src/postgkyl/diagnostics/discovery.py @@ -0,0 +1,71 @@ +"""Equation-blind output discovery -- Gkeyll's file-naming convention. + +The ONE home for "what outputs does this directory hold" (CLAUDE.md, +diagnostics layer). Every equation loader in ``gk/`` and every +program-scale diagnostic (layer 13) resolves files through here, never with +private ``glob`` logic of its own -- doctrine V, one home per fact. + +Ported from ``src_bak/postgkyl/loader.py``'s ``find_output_stems`` plus a new +``available_frames`` helper factored out of +``src_bak/postgkyl/gk/gk_quantities/gkquantity.py``'s ``_avail_frames_src`` +(the gyrokinetic quantity registry no longer globs on its own -- see +``diagnostics/gk/quantity.py``). +""" + +from __future__ import annotations + +import glob +import os + +from postgkyl import io + + +def find_output_stems(extensions: str = "gkyl", path: str = ".") -> dict: + """Map each extension to the sorted unique Gkeyll filename stems in ``path``. + + Frame indices and a trailing ``_restart`` are stripped from each stem by + :func:`postgkyl.io.parse_output_name` -- the one home for Gkeyll's naming + convention -- rather than by a private regex here. + + Args: + extensions: Comma-separated list of file extensions to scan. + path: Directory to scan. + + Returns: + A dict mapping each extension to a sorted list of unique stems. + """ + result = {} + for ext in extensions.split(","): + unique = [] + for fn in glob.glob(f"{path}/*.{ext:s}"): + stem = io.parse_output_name(os.path.basename(fn)).stem + if stem not in unique: + unique.append(stem) + result[ext] = sorted(unique) + return result + + +def available_frames(stem: str, *, frames: list[int] | None = None) -> set[int]: + """Set of available frame numbers for a ``.gkyl`` file family. + + Args: + stem: The file stem, including any trailing separator before the frame + number (e.g. ``"path/name-elc_M0_"``). + frames: Restrict the search to these candidate frame numbers instead of + globbing the whole directory (cheaper when the caller already has a + short candidate list). + + Returns: + The set of frame numbers for which ``.gkyl`` exists. + """ + found: set[int] = set() + if frames: + candidates = (f"{stem}{f}.gkyl" for f in frames + if os.path.isfile(f"{stem}{f}.gkyl")) + else: + candidates = glob.glob(f"{glob.escape(stem)}*.gkyl") + for f in candidates: + suffix = f[len(stem):-5] + if suffix.isdigit(): + found.add(int(suffix)) + return found diff --git a/src/postgkyl/diagnostics/gk/__init__.py b/src/postgkyl/diagnostics/gk/__init__.py new file mode 100644 index 00000000..aa528cf7 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/__init__.py @@ -0,0 +1,112 @@ +"""Gyrokinetic diagnostics: loading and equation-specific physics. + +The whole gyrokinetic-quantity stack -- naming-convention file resolution +(``quantity.py``), the derived-quantity physics (``quantities.py``), the +registry (``registry.py``), and the "physics-ready data by name" entry point +(``load_quantity.py``) -- lives together in this subpackage (see the layer-12 +instruction file's decision record): splitting resolution from physics would +give gyrokinetics two homes for one piece of equation knowledge. Only the +equation-blind stem/frame discovery is shared, via +``postgkyl.diagnostics.discovery``. Geometry-only transformations live below +this physics layer in ``postgkyl.operations.gyrokinetics``; the R-Z and +flux-surface names exported here are compatibility aliases. +""" + +from __future__ import annotations + +from .distf import load_distf, resolve_frames +from .load_quantity import available_quantities, load_quantity +from .quantities import ( + fetch_beta_from_bmag_press, + fetch_diamag_vel, + fetch_ExB_vel, + fetch_gradB_vel, + fetch_M1_from_H, + fetch_press_from_BiMax, + fetch_press_from_Max, + fetch_press_p, + fetch_Tpar_from_BiMax, + fetch_Tpar_from_M0_M1_M2par, + fetch_temp_from_Max, + fetch_temp_from_Tpar_Tperp, + fetch_Tperp_from_BiMax, + fetch_Tperp_from_M0_M2perp, +) +from .registry import gk_quant_registry + +# Layer 13: program-scale diagnostics ported from src_bak's apps/gk_*.py. +from .energy_balance import EnergyBalanceTraces, energy_balance_error, energy_balance +from .particle_balance import ParticleBalanceTraces, particle_balance, particle_balance_error +from .nodes import GKYL_GEOMETRY_ID, nodes, is_geo_mapc2p, multib_tag, nodes_to_RZ + +# Compatibility exports: canonical transformation APIs now live under +# postgkyl.operations.gyrokinetics. These imports are exact aliases. +from .rz import Geometry, RzProjection, gk_rz, map_to_rz, resolve_geometry, resolve_rz_projection +from .fluxsurf import FluxSurfaceGrid, extract_flux_surface, resolve_flux_surface_grid + +from typing import Annotated + +from postgkyl.cli_spec import ( + CommandSpec, + Execution, + KeyValue, + ResultPolicy, + Section, + command, +) + +_LOAD_SPEC = CommandSpec(Section.DIAGNOSTICS, Execution.LOAD) +_REPORT_SPEC = CommandSpec(Section.DIAGNOSTICS, + Execution.LOAD, + result=ResultPolicy.VALUE) +command(_LOAD_SPEC)(load_distf) +command(_LOAD_SPEC)(load_quantity) +energy_balance.__annotations__["bflux_files"] = Annotated[dict[str, str] | None, + KeyValue()] +particle_balance.__annotations__["bflux_files"] = Annotated[dict[str, str] + | None, + KeyValue()] +for _function in (energy_balance, particle_balance, nodes): + command(_REPORT_SPEC)(_function) + +__all__ = [ + "load_distf", + "resolve_frames", + "available_quantities", + "load_quantity", + "gk_quant_registry", + "fetch_beta_from_bmag_press", + "fetch_diamag_vel", + "fetch_ExB_vel", + "fetch_gradB_vel", + "fetch_M1_from_H", + "fetch_press_from_BiMax", + "fetch_press_from_Max", + "fetch_press_p", + "fetch_Tpar_from_BiMax", + "fetch_Tpar_from_M0_M1_M2par", + "fetch_temp_from_Max", + "fetch_temp_from_Tpar_Tperp", + "fetch_Tperp_from_BiMax", + "fetch_Tperp_from_M0_M2perp", + "EnergyBalanceTraces", + "energy_balance_error", + "energy_balance", + "ParticleBalanceTraces", + "particle_balance", + "particle_balance_error", + "GKYL_GEOMETRY_ID", + "nodes", + "is_geo_mapc2p", + "multib_tag", + "nodes_to_RZ", + "Geometry", + "RzProjection", + "gk_rz", + "map_to_rz", + "resolve_geometry", + "resolve_rz_projection", + "FluxSurfaceGrid", + "extract_flux_surface", + "resolve_flux_surface_grid", +] diff --git a/src/postgkyl/diagnostics/gk/distf.py b/src/postgkyl/diagnostics/gk/distf.py new file mode 100644 index 00000000..e414cf6b --- /dev/null +++ b/src/postgkyl/diagnostics/gk/distf.py @@ -0,0 +1,261 @@ +"""Loader for Gkeyll gyrokinetic distribution functions. + +Reads the saved ``Jf`` (distribution times one or more Jacobians) together +with the velocity/configuration Jacobians, divides them out, and +interpolates onto a nodal grid, optionally applying velocity- and +position-space coordinate mappings. + +Jf (phase-space) is weak-multiplied by jacobtot_inv (conf-space) via Gkeyll's +``gkyl_dg_mul_conf_phase_op_range`` staying gkyl-native and +already on the same (phase-space) grid as Jf, so a single ``interpolate()`` +at the end suffices; no separate jacobtot_inv interpolation or manual +NumPy reshape/broadcast is needed. The division by jacobvel, in contrast, +happens on the *raw* modal coefficient arrays via plain NumPy division on +the ``.values`` views, not Gkeyll's weak-divide kernel: ``jacobvel`` carries +no DG basis metadata of its own and is stored piecewise-constant per cell (a +single component), so Gkeyll's ``weak_div`` (which requires both operands' +component count to be a multiple of a shared basis's ``num_basis``) cannot +take it as an operand at all. Scaling every one of the coefficients by that +one per-cell constant is nonetheless the exact quotient, and commutes freely +with the weak conf x phase multiply above (both are linear in Jf's +coefficients), so the two can happen in either order. + +``resolve_frames``' range-discovery calls the shared +:mod:`postgkyl.diagnostics.discovery` helper instead of its own glob. +""" + +from __future__ import annotations + +from typing import Annotated + +from postgkyl import operations +from postgkyl.cli_spec import CliType +from postgkyl.gdata import GData, GDataGroup, load + +from .. import discovery + +FrameSpec = int | str | list[int] | tuple[int, ...] + + +def resolve_frames( + frame: FrameSpec, + *, + name: str, + species: str, + suffix: str = "", + block_idx: int | None = None, +) -> list[int]: + """Expand a frame specification into a concrete sorted list of frame indices. + + Args: + frame: An ``int`` (single frame); a ``list``/``tuple`` of ints; a string + with a single number (``"7"``) or comma-separated numbers + (``"0,2,4"``); or a ``'start:stop[:step]'`` / ``':'`` range (range + bounds default to the first/last frame discovered on disk). + name: Simulation name prefix. + species: Species name. + suffix: Distribution-file suffix (see :func:`load_distf`). + block_idx: Use block-specific files with a ``_b`` prefix. + + Returns: + A sorted list of concrete frame indices. + """ + if isinstance(frame, int): + return [frame] + if isinstance(frame, (list, tuple)): + return [int(f) for f in frame] + + frame_spec = str(frame).strip() + if "," in frame_spec: + return [int(f.strip()) for f in frame_spec.split(",")] + if ":" not in frame_spec: + return [int(frame_spec)] + + prefix = f"{name}_b{block_idx}" if block_idx is not None else name + frame_infix = f"{suffix}_" if suffix else "" + stem = f"{prefix}-{species}_{frame_infix}" + available = sorted(discovery.available_frames(stem)) + if not available: + raise ValueError( + f"No distribution frames found matching '{stem}.gkyl'.") + parts = frame_spec.split(":") + if len(parts) > 3: + raise ValueError( + f"Invalid frame range {frame_spec!r}; expected start:stop[:step].") + lower = int(parts[0]) if parts[0] else available[0] + upper = int(parts[1]) if parts[1] else available[-1] + 1 + step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 + if step <= 0: + raise ValueError("Frame range step must be a positive integer.") + resolved = [ + f for f in available if lower <= f < upper and (f - lower) % step == 0 + ] + if not resolved: + raise ValueError( + f"Frame range {frame_spec!r} matches no files for '{stem}.gkyl'." + ) + return resolved + + +def load_distf( + name: str, + species: str, + frame: Annotated[FrameSpec, CliType(str)], + *, + tag: str = "f", + suffix: str = "", + use_c2p_vel: bool = False, + use_mc2nu: bool = False, + use_mapc2p: bool = False, + block_idx: int | None = None, + num_interp: int | None = None, + jf_file: str | None = None, + mapc2p_vel_file: str | None = None, + jacobvel_file: str | None = None, + mc2nu_file: str | None = None, + mapc2p_file: str | None = None, + jacobtot_inv_file: str | None = None, +) -> GData | GDataGroup: + """Build real distribution functions from saved ``Jf`` data. + + A scalar frame returns one :class:`~postgkyl.gdata.gdata.GData`. A list, + tuple, comma-separated string, or range returns a + :class:`~postgkyl.gdata.gdatagroup.GDataGroup`, whose fluent operations + broadcast over the loaded frames. + + Args: + name: Simulation name prefix. + species: Species name. + frame: Frame index, comma-separated indices, or a + ``start:stop[:step]`` range; ``:`` selects every available frame. + tag: Tag for the resulting dataset. + suffix: Use ``-__.gkyl`` as the input. + use_c2p_vel: Convert velocity-space computational coordinates to + physical ones using the ``mapc2p_vel`` mapping. + use_mc2nu: Convert non-uniform computational coordinates to + field-aligned ones. + use_mapc2p: Convert position-space computational coordinates to + Cartesian/cylindrical. + block_idx: Use block-specific files with a ``_b`` prefix. + num_interp: Interpolate onto a general mesh of the specified amount + (default: ``poly_order + 1`` points per cell). + jf_file: Explicit saved-distribution filename override. + mapc2p_vel_file: Explicit velocity-coordinate mapping filename override. + jacobvel_file: Explicit velocity-space Jacobian filename override. + mc2nu_file: Explicit field-aligned coordinate mapping filename override. + mapc2p_file: Explicit configuration-space mapping filename override. + jacobtot_inv_file: Explicit inverse total-Jacobian filename override. + + Returns: + One interpolated distribution function for a scalar frame, or a fluent + group holding one distribution function per requested frame. + """ + frames = resolve_frames(frame, + name=name, + species=species, + suffix=suffix, + block_idx=block_idx) + datasets = [ + _load_distf_frame(name=name, + species=species, + frame=resolved_frame, + tag=tag, + suffix=suffix, + use_c2p_vel=use_c2p_vel, + use_mc2nu=use_mc2nu, + use_mapc2p=use_mapc2p, + block_idx=block_idx, + num_interp=num_interp, + jf_file=jf_file, + mapc2p_vel_file=mapc2p_vel_file, + jacobvel_file=jacobvel_file, + mc2nu_file=mc2nu_file, + mapc2p_file=mapc2p_file, + jacobtot_inv_file=jacobtot_inv_file) + for resolved_frame in frames + ] + + is_series = isinstance(frame, + (list, tuple)) or (isinstance(frame, str) and + ("," in frame or ":" in frame)) + if not is_series: + return datasets[0] + for dataset, resolved_frame in zip(datasets, frames): + dataset.set_label(str(resolved_frame)) + return GDataGroup(datasets) + + +def _load_distf_frame( + name: str, + species: str, + frame: int, + *, + tag: str = "f", + suffix: str = "", + use_c2p_vel: bool = False, + use_mc2nu: bool = False, + use_mapc2p: bool = False, + block_idx: int | None = None, + num_interp: int | None = None, + jf_file: str | None = None, + mapc2p_vel_file: str | None = None, + jacobvel_file: str | None = None, + mc2nu_file: str | None = None, + mapc2p_file: str | None = None, + jacobtot_inv_file: str | None = None, +) -> GData: + """Load and transform one resolved distribution-function frame.""" + prefix = f"{name}_b{block_idx}" if block_idx is not None else name + frame_infix = f"{suffix}_" if suffix else "" + + if jf_file is None: + jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" + if mapc2p_vel_file is None: + mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" + if jacobvel_file is None: + jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" + if mc2nu_file is None: + mc2nu_file = f"{prefix}-geo_corn_mc2nu_pos_deflated.gkyl" + if mapc2p_file is None: + mapc2p_file = f"{prefix}-geo_corn_mapc2p_deflated.gkyl" + if jacobtot_inv_file is None: + jacobtot_inv_file = f"{prefix}-geo_int_jacobtot_inv.gkyl" + + jf_data = load(jf_file) + # jacobvel is stored piecewise-constant per cell (see module docstring): one + # coefficient per cell is exactly a poly_order=0 DG field, so this is real + # metadata, not a guess -- it silences the load-time "missing basis" + # warning honestly instead of leaving it to fire on every distf load. + jacobvel_data = load(jacobvel_file, basis_type="serendipity", poly_order=0) + jacobtot_inv_data = load(jacobtot_inv_file) + + weak_product = jf_data * jacobtot_inv_data + f_coeffs = weak_product.values / jacobvel_data.values + f_modal = weak_product._result(weak_product.grid, f_coeffs) + # The composed distribution's true basis (gkhybrid, p1) is fixed by this + # diagnostic's convention, not necessarily what jf_data's own file header + # implies -- basis_type/poly_order/value_form are load-time-fixed + # properties, so the override lands on ctx here rather than as an + # interpolate() argument. + f_modal.ctx.update(basis_type="gkhybrid", poly_order=1, value_form="modal") + + interpolated = f_modal.interpolate(num_interp=num_interp) + out = interpolated._result(interpolated.grid, interpolated.values, tag=tag) + + # Coordinate maps run on the already-interpolated data via the shared map + # verb. Velocity space (c2p_vel) deforms the trailing axes; configuration + # space (mc2nu / mapc2p) deforms the leading ones. + grid_type = [] + if use_c2p_vel: + mc2p_vel = load(mapc2p_vel_file, poly_order=1, basis_type="serendipity") + out = operations.map(out, mc2p_vel, space="vel") + grid_type.append("c2p_vel") + if use_mc2nu: + out = operations.map(out, mc2nu_file, space="conf") + grid_type.append("mc2nu") + elif use_mapc2p: + out = operations.map(out, mapc2p_file, space="conf") + grid_type.append("mapc2p") + if grid_type: + out.ctx["grid_type"] = " + ".join(grid_type) + return out diff --git a/src/postgkyl/diagnostics/gk/energy_balance.py b/src/postgkyl/diagnostics/gk/energy_balance.py new file mode 100644 index 00000000..1e384f04 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/energy_balance.py @@ -0,0 +1,390 @@ +"""Gyrokinetic energy-balance diagnostic. + +Ported from ``src_bak/postgkyl/apps/gk_energy_balance.py``. Reads the +integrated time-trace files a gyrokinetic simulation writes (field/apar +energy rate of change, integrated Hamiltonian moments of ``df/dt``, of the +source(s), and of the boundary particle fluxes), sums them over species and +(for multiblock runs) blocks, and plots the energy-balance residual:: + + E_err = S - bflux - (df/dt - dfield/dt [- dapar/dt]) + +Typer options become explicit keyword-only parameters; the old CLI's dataset +stack (``ctx.obj.data``) and ``verb_print`` echo are dropped -- the computed +traces come back as an :class:`EnergyBalanceTraces` alongside the Figure. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import matplotlib.pyplot as plt +import numpy as np + +from . import utils + +_DIRS = ("x", "y", "z") +_EDGES = ("lower", "upper") +_LINE_STYLES = ("-", "--", ":", "-.") +_XY_LABEL_FONT_SIZE = 17 +_TITLE_FONT_SIZE = 17 +_TICK_FONT_SIZE = 14 +_LEGEND_FONT_SIZE = 14 + +# Hamiltonian-moments files store (M0, M1, M2) per component; energy balance +# uses the M2 (Hamiltonian/energy) moment, index 2. +_ENERGY_MOMENT = 2 + + +@dataclass(frozen=True) +class EnergyBalanceTraces: + """Computed energy-balance time traces (all 1-D, aligned to ``time``). + + Attributes: + time: Time stamps of the (dominant) ``fdot`` trace. + fdot: Rate of change of the Hamiltonian moment of the distribution + function, summed over species and blocks. + src: Rate of change from sources, or ``None`` if no source file was + found for any species/block. + bflux_tot: Rate of change from boundary particle fluxes, or ``None`` if + none were found. + field_dot: Rate of change of the field energy. + apar_dot: Rate of change of the vector-potential energy (electromagnetic + simulations only), or ``None``. + mom_err: The energy-balance residual (``None`` when ``relative_error``). + mom_err_norm: The *relative* energy-balance residual (only set when + ``relative_error=True``; ``None`` otherwise). + """ + + time: np.ndarray + fdot: np.ndarray + src: np.ndarray | None + bflux_tot: np.ndarray | None + field_dot: np.ndarray + apar_dot: np.ndarray | None + mom_err: np.ndarray | None + mom_err_norm: np.ndarray | None = None + + +def _accumulate(target: np.ndarray | None, addend) -> np.ndarray: + """Sum ``addend`` into ``target`` (over species/blocks), copying on first + use so the caller's array is never mutated in place.""" + addend = np.asarray(addend) + return addend.copy() if target is None else target + addend + + +def energy_balance_error(fdot: np.ndarray, + src: np.ndarray, + bflux_tot: np.ndarray, + field_dot: np.ndarray, + apar_dot: np.ndarray | None = None) -> np.ndarray: + """The energy-balance residual: ``S - bflux - (df/dt - dfield/dt [- dapar/dt])``. + + Pure array arithmetic -- the one formula every energy-balance trace + (single- or multi-block, single- or multi-species) reduces to once the + per-species/per-block sums are in hand. + """ + fdot_terms = fdot - field_dot + if apar_dot is not None: + fdot_terms = fdot_terms - apar_dot + return src - bflux_tot - fdot_terms + + +def _block_prefix(file_prefix: str, block_idx: int) -> str: + return file_prefix.replace("*", str(block_idx)) + + +def _resolve(path: str, + override: str | None, + default: str, + block_idx: int, + species: str | None = None) -> str: + """Resolve a file-family member's path: ``override`` (with ``*`` + substituted for the block index, then the species) if given, else the + naming-convention ``default``.""" + if override is None: + return default + resolved = (path + override).replace("*", str(block_idx), 1) + if species is not None: + resolved = resolved.replace("*", species) + return resolved + + +def energy_balance( + name: str, + species: list[str], + *, + path: str = "./", + relative_error: bool = False, + multib: str = "-10", + field_dot_file: str | None = None, + apar_dot_file: str | None = None, + fdot_file: str | None = None, + source_file: str | None = None, + bflux_files: dict[str, str] | None = None, + f_file: str | None = None, + field_file: str | None = None, + apar_file: str | None = None, + dt_file: str | None = None, + logy: bool = False, + absy: bool = False, + xlabel: str = "Time (s)", + ylabel: str | None = None, + title: str | None = None, + indent_left: float = 0.0, + add_width: float = 0.0, + show: bool = False, + saveas: str | None = None, +) -> tuple[plt.Figure, EnergyBalanceTraces]: + """Plot (and compute) the energy balance of a gyrokinetic simulation. + + Requires, per species (named ``-``): an + ``_fdot_integrated_moms.gkyl`` file, and (only if the run had sources or + non-periodic boundaries) ``_source_integrated_moms.gkyl`` and + ``_bflux__integrated_HamiltonianMoments.gkyl`` files. A + ``-field_energy_dot.gkyl`` file is required; ``-apar_energy_dot + .gkyl`` is read if present (electromagnetic simulations). If + ``relative_error`` is requested, the corresponding non-``_dot`` + (``_integrated_moms.gkyl``/``field_energy.gkyl``/``apar_energy.gkyl``) and + ``-dt.gkyl`` files are also required. + + Args: + name: Simulation name (also the file prefix). + species: Species names to sum over. + path: Directory holding the simulation output. + relative_error: Plot the relative error instead of every balance term. + multib: ``"-10"`` (default) for a single block; ``"-1"`` to discover + every block; otherwise a comma list or ``'start:stop[:step]'`` slice + of block indices (see :func:`~postgkyl.diagnostics.gk.utils. + get_block_indices`). + field_dot_file: Explicit field-energy derivative path override. + apar_dot_file: Explicit parallel-vector-potential energy derivative path. + fdot_file: Explicit distribution derivative moments path override. + source_file: Explicit source moments path override. + bflux_files: Optional per-boundary path overrides, keyed by + ``""`` (e.g. ``"xlower"``); unlisted boundaries use + the naming convention. + f_file: Explicit integrated distribution moments path override. + field_file: Explicit field-energy path override. + apar_file: Explicit parallel-vector-potential energy path override. + dt_file: Explicit time-step history path override. In path overrides, + ``*`` stands for block index and, for per-species files, species name. + logy: Log-scale the y axis. + absy: Take the absolute value of every trace before plotting. + xlabel: Horizontal-axis label. + ylabel: Vertical-axis label; use the derived formula when ``None``. + title: Figure title; use the balance description when ``None``. + indent_left: Horizontal axes-position adjustment in figure units. + add_width: Axes-width adjustment in figure units. + show: Call ``plt.show()`` before returning. + saveas: If given, save the figure to this path. + + Returns: + ``(figure, traces)``. + + Raises: + FileNotFoundError: if a required file family is missing. + """ + path = path.rstrip("/") + "/" + bflux_files = bflux_files or {} + + file_prefix = f"{path}{name}-" if multib == "-10" else f"{path}{name}_b*-" + probe = fdot_file or (file_prefix + species[0] + "_fdot_integrated_moms.gkyl") + blocks = utils.get_block_indices(multib, probe) + + fig = plt.figure(figsize=(7.5, 4.5)) + ax = fig.add_axes([0.11 + indent_left, 0.15, 0.87 + add_width, 0.78]) + ax.plot([-1.0, 1.0], [0.0, 0.0], color="grey", linestyle=":", linewidth=1) + + absy_func = np.abs if absy else (lambda v: v) + + field_dot = apar_dot = fdot = src = bflux_tot = None + has_apar_dot = has_src = has_bflux = False + time_fdot = time_field_dot = time_apar_dot = time_bflux_tot = None + + for block_idx in blocks: + block_prefix = _block_prefix(file_prefix, block_idx) + + fd_name = _resolve(path, field_dot_file, + block_prefix + "field_energy_dot.gkyl", block_idx) + found, t, v, _ = utils.read_time_trace_if_present(fd_name) + if not found: + raise FileNotFoundError(f"Required file not found: {fd_name}") + time_field_dot, field_dot_pb = t, v + + ad_name = _resolve(path, apar_dot_file, + block_prefix + "apar_energy_dot.gkyl", block_idx) + has_apar_dot, t, v, _ = utils.read_time_trace_if_present(ad_name) + if has_apar_dot: + time_apar_dot, apar_dot_pb = t, v + + fdot_pb = src_pb = bflux_tot_pb = None + for sp in species: + fdot_name = _resolve(path, fdot_file, + block_prefix + sp + "_fdot_integrated_moms.gkyl", + block_idx, sp) + found, t, v, _ = utils.read_time_trace_if_present(fdot_name) + if not found: + raise FileNotFoundError(f"Required file not found: {fdot_name}") + time_fdot = t + fdot_sp = v[:, _ENERGY_MOMENT] + + src_name = _resolve(path, source_file, + block_prefix + sp + "_source_integrated_moms.gkyl", + block_idx, sp) + has_src, t, v, _ = utils.read_time_trace_if_present(src_name) + if has_src: + src_sp = v[:, _ENERGY_MOMENT] + else: + src_sp = 0.0 * fdot_sp + + bflux_terms = [] + for d in _DIRS: + for e in _EDGES: + key = d + e + bf_name = _resolve( + path, bflux_files.get(key), block_prefix + sp + + f"_bflux_{d}{e}_integrated_HamiltonianMoments.gkyl", block_idx, + sp) + found_b, t, v, _ = utils.read_time_trace_if_present(bf_name) + if found_b: + has_bflux = True + time_bflux_tot = t + bflux_terms.append(v[:, _ENERGY_MOMENT]) + bflux_sp = sum(bflux_terms) if bflux_terms else 0.0 * fdot_sp + + fdot_pb = _accumulate(fdot_pb, fdot_sp) + src_pb = _accumulate(src_pb, src_sp) + bflux_tot_pb = _accumulate(bflux_tot_pb, bflux_sp) + + field_dot = _accumulate(field_dot, field_dot_pb) + if has_apar_dot: + apar_dot = _accumulate(apar_dot, apar_dot_pb) + fdot = _accumulate(fdot, fdot_pb) + src = _accumulate(src, src_pb) + bflux_tot = _accumulate(bflux_tot, bflux_tot_pb) + + legend_handles = [] + legend_strings = [] + + if not relative_error: + src = src.copy() + src[0] = 0.0 # No fdot/bflux contribution at t=0. + + mom_err = energy_balance_error(fdot, src, bflux_tot, field_dot, + apar_dot if has_apar_dot else None) + + if has_src: + h, = ax.plot(time_fdot, absy_func(src), linestyle=_LINE_STYLES[2]) + legend_handles.append(h) + legend_strings.append(r"$\mathcal{S}$") + if has_bflux: + h, = ax.plot(time_bflux_tot, + absy_func(-bflux_tot), + linestyle=_LINE_STYLES[1]) + legend_handles.append(h) + legend_strings.append( + r"$-\int_{\partial \Omega}\mathrm{d}\mathbf{S}\cdot\mathbf{\dot{R}}f$" + ) + h, = ax.plot(time_field_dot, + absy_func(-field_dot), + linestyle=":", + marker="+", + markevery=8) + legend_handles.append(h) + legend_strings.append(r"$-\dot{\phi}$") + if has_apar_dot: + h, = ax.plot(time_apar_dot, + absy_func(-apar_dot), + linestyle=":", + marker="+", + markevery=8) + legend_handles.append(h) + legend_strings.append(r"$-\dot{A}_{\parallel}$") + h, = ax.plot(time_fdot, absy_func(-fdot), linestyle=_LINE_STYLES[0]) + legend_handles.append(h) + legend_strings.append(r"$-\dot{f}$") + h, = ax.plot(time_fdot, absy_func(mom_err), linestyle=_LINE_STYLES[3]) + legend_handles.append(h) + legend_strings.append(r"$E_{\dot{\mathcal{E}}}=$" + "".join(legend_strings)) + + ax.legend(legend_handles, + legend_strings, + fontsize=_LEGEND_FONT_SIZE, + frameon=False) + + ylabel_string = ylabel or "" + title_string = title or r"Energy balance" + mom_err_norm = None + else: + dt_name = _resolve(path, dt_file, + file_prefix.replace("_b*", "") + "dt.gkyl", 0) + _, time_dt, dt, _ = utils.read_time_trace_if_present(dt_name) + + field = apar = distf = None + for block_idx in blocks: + block_prefix = _block_prefix(file_prefix, block_idx) + + fld_name = _resolve(path, field_file, block_prefix + "field_energy.gkyl", + block_idx) + has_field, t, v, _ = utils.read_time_trace_if_present(fld_name) + field_pb = v if has_field else None + + ap_name = _resolve(path, apar_file, block_prefix + "apar_energy.gkyl", + block_idx) + has_apar, t, v, _ = utils.read_time_trace_if_present(ap_name) + apar_pb = v if has_apar else None + + distf_pb = None + for sp in species: + f_name = _resolve(path, f_file, + block_prefix + sp + "_integrated_moms.gkyl", + block_idx, sp) + _, t, v, _ = utils.read_time_trace_if_present(f_name) + distf_pb = _accumulate(distf_pb, v[:, _ENERGY_MOMENT]) + + field = _accumulate(field, field_pb) + if has_apar: + apar = _accumulate(apar, apar_pb) + distf = _accumulate(distf, distf_pb) + + field, field_dot = field[1:], field_dot[1:] + if has_apar: + apar, apar_dot = apar[1:], apar_dot[1:] + fdot, src, bflux_tot, distf = fdot[1:], src[1:], bflux_tot[1:], distf[1:] + + mom_err = energy_balance_error(fdot, src, bflux_tot, field_dot, + apar_dot if has_apar else None) + denom = (distf - field - apar) if has_apar else (distf - field) + mom_err_norm = mom_err * dt / denom + + ax.plot(time_dt, absy_func(mom_err_norm)) + + ylabel_string = ylabel or r"$E_{\dot{\mathcal{E}}}~\Delta t/\mathcal{E}$" + title_string = title or r"Relative error in energy conservation" + mom_err = None + + if logy: + ax.set_yscale("log") + if absy and ylabel_string: + ylabel_string = r"|" + ylabel_string + r"|" + + ax.set_xlabel(xlabel, fontsize=_XY_LABEL_FONT_SIZE) + ax.set_ylabel(ylabel_string, fontsize=_XY_LABEL_FONT_SIZE) + ax.set_title(title_string, fontsize=_TITLE_FONT_SIZE) + ax.set_xlim(time_fdot[0], time_fdot[-1]) + utils.set_tick_font_size(ax, _TICK_FONT_SIZE) + + if saveas: + fig.savefig(saveas) + if show: + plt.show() + + traces = EnergyBalanceTraces(time=time_fdot, + fdot=fdot, + src=src if has_src else None, + bflux_tot=bflux_tot if has_bflux else None, + field_dot=field_dot, + apar_dot=apar_dot if has_apar_dot else None, + mom_err=mom_err, + mom_err_norm=mom_err_norm) + return fig, traces diff --git a/src/postgkyl/diagnostics/gk/fluxsurf.py b/src/postgkyl/diagnostics/gk/fluxsurf.py new file mode 100644 index 00000000..64bcc7bc --- /dev/null +++ b/src/postgkyl/diagnostics/gk/fluxsurf.py @@ -0,0 +1,23 @@ +"""Compatibility aliases for gyrokinetic flux-surface operations. + +Canonical imports live in :mod:`postgkyl.operations.gyrokinetics`; this path +is scheduled for removal in the next major version. +""" + +from postgkyl.operations.gyrokinetics.fluxsurf import ( + FluxSurfaceGrid, + Geometry, + extract_flux_surface, + flux_surface_grids, + grid_for, + resolve_flux_surface_grid, +) + +__all__ = [ + "Geometry", + "FluxSurfaceGrid", + "extract_flux_surface", + "flux_surface_grids", + "grid_for", + "resolve_flux_surface_grid", +] diff --git a/src/postgkyl/diagnostics/gk/load_quantity.py b/src/postgkyl/diagnostics/gk/load_quantity.py new file mode 100644 index 00000000..20a09a61 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/load_quantity.py @@ -0,0 +1,140 @@ +"""Loader for pre-named gyrokinetic quantities. + +Resolves a quantity name through the :mod:`postgkyl.diagnostics.gk. +registry`, loads the required source files, computes the quantity, and +returns ready datasets. Ported from +``src_bak/postgkyl/loaders/gk_quantity.py``. +""" + +from __future__ import annotations + +from typing import Annotated, TYPE_CHECKING + +from postgkyl.cli_spec import ChoiceProvider, KeyValue +from .registry import gk_quant_registry + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def available_quantities() -> list[str]: + """Return the sorted list of registered quantity names.""" + return gk_quant_registry.list() + + +def load_quantity( + quantity: Annotated[str, ChoiceProvider(available_quantities)], + species: str | None, + name: str, + frame: str | None = None, + *, + path: str = "./", + tag: str = "default", + label: str | None = None, + direction: int | None = None, + mass: float | None = None, + charge: float | None = None, + gamma_e: float | None = None, + gamma_i: float | None = None, + kind: str | None = None, + read_options: Annotated[dict[str, str] | None, + KeyValue()] = None, +) -> list: + """Load and compute a pre-named gyrokinetic quantity. + + Args: + quantity: Registered quantity name (see :func:`available_quantities`). + species: Species name, or a comma-separated list of them; ``None`` for + species-independent quantities. + name: Simulation name prefix (e.g. ``'gk_sheath_2x2v_p1'``). + frame: Frame number, comma-separated list, or ``'start:stop[:step]'`` + range; ``':'``/``None`` selects all available frames. + path: Directory containing the simulation files. + tag: Tag for the output dataset(s); suffixed with the species when more + than one species is requested. + label: Label override; defaults to the quantity's registered label. + direction: Vector direction for quantities that expose components. + mass: Species mass used by quantities that require it. + charge: Species charge used by quantities that require it. + gamma_e: Electron adiabatic index for sound-speed quantities. + gamma_i: Ion adiabatic index for sound-speed quantities. + kind: Named variant accepted by a quantity provider. + read_options: Additional provider options as repeated key/value entries. + + Returns: + A list of computed ``GDataState`` datasets. + + Raises: + ValueError: if ``quantity`` is not registered, or it is an + ``is_multi_species`` quantity requested without a species list. + """ + extra = dict(read_options or {}) + for key, value in (("dir", direction), ("mass", mass), ("charge", charge), + ("gamma_e", gamma_e), ("gamma_i", gamma_i), ("kind", + kind)): + if value is not None: + extra[key] = value + + if not gk_quant_registry.has(quantity): + valid = gk_quant_registry.list() + raise ValueError(f"Unknown quantity '{quantity}'. Available quantities: " + f"{', '.join(valid)}.") + + gkquant = gk_quant_registry.get(quantity) + path = path.rstrip("/") + "/" + species_list = [s.strip() for s in species.split(",")] if species else [None] + + frame_inp = str(frame) if frame is not None else None + + if gkquant.is_multi_species: + # Combine every species into a single dataset (e.g. the sound speed), + # so it is fetched once for the whole species list instead of once + # per species. + if species_list == [None]: + raise ValueError( + f"Quantity '{quantity}' combines several species, so it needs a " + "species list, e.g. --species elc,ion.") + + src_combo_idx, frames = gkquant.get_avail_source_multi( + path, name, species_list, frame_inp) + + datasets: list["GDataState"] = [] + for fr in frames: + out = gkquant.fetch_multi(path, name, species_list, fr, src_combo_idx, + **extra) + + out_label = label if label is not None else gkquant.get_label() + if len(frames) > 1: + out_label += f" f{fr}" + out.set_label(out_label) + out.set_tag(tag) + + datasets.append(out) + return datasets + + datasets: list["GDataState"] = [] + for species_idx, sp in enumerate(species_list): + src_combo_idx, frames = gkquant.get_avail_source(path, name, sp, frame_inp) + + # Tells the fetch functions which entry of a per-species '--extra' array + # (e.g. 'mass=1,2,3') applies to the species being computed. + species_extra = dict(extra, species_idx=species_idx) + + for fr in frames: + out = gkquant.fetch(path, name, sp, fr, src_combo_idx, **species_extra) + + default_label = gkquant.get_label(species=sp, direction=extra.get("dir")) + if label is not None: + out_label = label + (f" {sp}" if len(species_list) > 1 else "") + else: + out_label = default_label + if len(frames) > 1: + out_label += f" f{fr}" + out.set_label(out_label) + + out_tag = tag + (f"_{sp}" if len(species_list) > 1 else "") + out.set_tag(out_tag) + + datasets.append(out) + + return datasets diff --git a/src/postgkyl/diagnostics/gk/nodes.py b/src/postgkyl/diagnostics/gk/nodes.py new file mode 100644 index 00000000..bed00ff5 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/nodes.py @@ -0,0 +1,260 @@ +"""Gyrokinetic grid-node diagnostic. + +Ported from ``src_bak/postgkyl/apps/gk_nodes.py``: plots the nodes of a +(possibly multiblock, possibly mapc2p) grid, connected by their cell edges, +with an optional overlay of the poloidal-flux contours/colormap and a vacuum- +vessel wall outline. +""" + +from __future__ import annotations + +from itertools import cycle + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.collections import LineCollection + +# GKYL_GEOMETRY_ID remains an intentional diagnostics.gk compatibility export. +from postgkyl.operations.gyrokinetics.geometry import ( # noqa: F401 + GKYL_GEOMETRY_ID, is_geo_mapc2p, +) + +from . import utils + +_XY_LABEL_FONT_SIZE = 17 +_TITLE_FONT_SIZE = 17 +_TICK_FONT_SIZE = 14 +_COLORBAR_LABEL_FONT_SIZE = 14 + + +def nodes_to_RZ(nodes: np.ndarray, + is_mapc2p: bool) -> tuple[np.ndarray, np.ndarray]: + """Compute the major-radius/vertical-location (R, Z) variables from a + grid-nodes array. + + Args: + nodes: Node coordinates, shape ``(*cell_shape, 3)`` holding Cartesian + (X, Y, Z) for mapc2p geometry, or ``(*cell_shape, 2+)`` holding + (R, Z, [phi]) otherwise. A size-1 ``y`` axis is sliced out (at index + 0) for 3-D cell shapes. + is_mapc2p: Whether ``nodes`` holds Cartesian coordinates (True) or + already (R, Z, ...) coordinates (False). + + Returns: + ``(majorR, vertZ)``. + """ + yidx = 0 # Index in the y direction to slice 3-D node arrays at. + + nx_nod = np.shape(nodes) + cdim = np.size(nx_nod) - 1 + cart_dim = 3 + + lo_idx = [[0 for _ in range(cdim)] + [cd] for cd in range(cart_dim)] + up_idx = [[nx_nod[d] for d in range(cdim)] + [cd + 1] + for cd in range(cart_dim)] + + if cdim == 3: + for cd in range(cart_dim): + lo_idx[cd][1] = yidx + up_idx[cd][1] = yidx + 1 + + slices = [[slice(lo_idx[cd][d], up_idx[cd][d]) for d in range(cdim + 1)] + for cd in range(cart_dim)] + + if is_mapc2p: + cart_x = [np.squeeze(nodes[tuple(slices[d])]) for d in range(cart_dim)] + major_r = np.sqrt(np.power(cart_x[0], 2) + np.power(cart_x[1], 2)) + vert_z = cart_x[2] + else: + major_r = np.squeeze(nodes[tuple(slices[0])]) + vert_z = np.squeeze(nodes[tuple(slices[1])]) + + return major_r, vert_z + + +def multib_tag(base: str, block_idx: int, num_blocks: int) -> str: + """Tag a per-block artifact, adding a ``_b`` suffix only when there + is more than one block.""" + return f"{base}_b{block_idx}" if num_blocks > 1 else base + + +def _parse_levels(clevels: str | None, cnlevels: int) -> np.ndarray | int: + if clevels is None: + return cnlevels + if ":" in clevels: + s = clevels.split(":") + return np.linspace(float(s[0]), float(s[1]), int(s[2])) + return np.array([float(v) for v in clevels.split(",") if v]) + + +def nodes( + name: str, + *, + path: str = "./", + multib: str = "-10", + nodes_file: str | None = None, + psi_file: str | None = None, + wall_file: str | None = None, + contour: bool = False, + clevels: str | None = None, + cnlevels: int = 11, + fixaspect: bool = False, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + xlabel: str = "R (m)", + ylabel: str = "Z (m)", + zlabel: str = r"$\psi$", + title: str | None = None, + indent_left: float = 0.0, + add_width: float = 0.0, + multib_unicolor: bool = False, + show: bool = False, + saveas: str | None = None, +) -> plt.Figure: + """Plot the nodes of a (possibly multiblock) grid, with optional overlays. + + Args: + name: Simulation name (also the file prefix). + path: Directory holding the simulation output. + multib: ``"-10"`` (default) for a single block; ``"-1"`` to discover + every block; otherwise a comma list or ``'start:stop[:step]'`` slice + of block indices. + nodes_file: Override for the ``-nodes.gkyl`` grid-nodes file + (``*`` stands for the block index); an absolute path is used as-is. + psi_file: Optional poloidal-flux file to overlay (interpolated p2 tensor + basis); an absolute path is used as-is. + wall_file: Optional CSV ``(R, Z)`` vacuum-vessel wall outline to overlay. + contour: Draw ``psi_file`` as contour lines instead of a colormesh. + clevels: Contour levels: comma-separated values, or a + ``'start:stop:nlevels'`` range; defaults to ``cnlevels`` automatic + levels when ``None``. + cnlevels: Number of automatic contour levels (ignored if ``clevels`` is + given). + fixaspect: Enforce equal R/Z scaling (unused placeholder kept for + interface symmetry with the old CLI's ``--fix_aspect``; the figure is + already built to the data's aspect ratio). + xlim: Optional horizontal-axis limits. + ylim: Optional vertical-axis limits. + xlabel: Horizontal-axis label. + ylabel: Vertical-axis label. + zlabel: Poloidal-flux colorbar label. + title: Figure title. + indent_left: Horizontal axes-position adjustment in figure units. + add_width: Axes-width adjustment in figure units. + multib_unicolor: Use one color for every block instead of cycling. + show: Call ``plt.show()`` before returning. + saveas: If given, save the figure to this path. + + Returns: + The populated Figure. + """ + path = path.rstrip("/") + "/" + file_prefix = f"{path}{name}-" if multib == "-10" else f"{path}{name}_b*-" + + if nodes_file: + resolved_nodes_file = nodes_file if nodes_file[ + 0] == "/" else path + nodes_file + else: + resolved_nodes_file = file_prefix + "nodes.gkyl" + + blocks = utils.get_block_indices(multib, resolved_nodes_file) + + major_r_ex = [1e9, -1e9] + vert_z_ex = [1e9, -1e9] + block_nodes = {} + for block_idx in blocks: + grid, nodes, gdat = utils.read_gfile( + resolved_nodes_file.replace("*", str(block_idx))) + mapc2p = is_geo_mapc2p(gdat.ctx) + major_r, vert_z = nodes_to_RZ(nodes, mapc2p) + block_nodes[block_idx] = (major_r, vert_z, gdat) + major_r_ex = [ + min(major_r_ex[0], np.amin(major_r)), + max(major_r_ex[1], np.amax(major_r)) + ] + vert_z_ex = [ + min(vert_z_ex[0], np.amin(vert_z)), + max(vert_z_ex[1], np.amax(vert_z)) + ] + + length_r = major_r_ex[1] - major_r_ex[0] + length_z = vert_z_ex[1] - vert_z_ex[0] + aspect_ratio = length_r / length_z + + ax_pos = [ + 0.82 - (8.36 * aspect_ratio) / (8.36 * aspect_ratio + 2.5) + indent_left, + 0.08, (8.36 * aspect_ratio) / (8.36 * aspect_ratio + 2.5) + add_width, + 0.88 + ] + cax_pos = [ax_pos[0] + ax_pos[2] + 0.01, ax_pos[1], 0.02, ax_pos[3]] + fig = plt.figure(figsize=(8.36 * aspect_ratio + 2.5, 8.36 + 1.14)) + ax = fig.add_axes(ax_pos) + + color_list = plt.rcParams["axes.prop_cycle"].by_key()["color"] + block_colors = cycle([color_list[0]] if multib_unicolor else color_list) + + for block_idx in blocks: + major_r, vert_z, gdat = block_nodes[block_idx] + ax.plot(major_r, vert_z, marker=".", color="k", linestyle="none") + + cell_color = next(block_colors) + if major_r.ndim <= 1: + ax.plot(major_r, vert_z, color=cell_color, linestyle="-") + else: + segs_constx = np.stack((major_r, vert_z), axis=2) + segs_consty = segs_constx.transpose(1, 0, 2) + ax.add_collection(LineCollection(segs_constx, color=cell_color)) + ax.add_collection(LineCollection(segs_consty, color=cell_color)) + + colorbar = True + if psi_file: + resolved_psi = psi_file if psi_file[0] == "/" else path + psi_file + psi_grid, psi_values, _ = utils.read_interpolated_gfile(resolved_psi, + poly_order=2, + basis_type="tensor") + psi_grid_cc = [ + 0.5 * (psi_grid[d][:-1] + psi_grid[d][1:]) for d in range(len(psi_grid)) + ] + + levels = _parse_levels(clevels, cnlevels) + if isinstance(levels, np.ndarray) and levels.size == 1: + colorbar = False + + if contour: + im = ax.contour(psi_grid_cc[0], psi_grid_cc[1], psi_values.transpose(), + levels) + else: + im = ax.pcolormesh(psi_grid[0], + psi_grid[1], + psi_values.transpose(), + cmap="inferno") + + if colorbar: + cbar_ax = fig.add_axes(cax_pos) + cbar = fig.colorbar(im, ax=ax, cax=cbar_ax) + cbar.ax.tick_params(labelsize=_TICK_FONT_SIZE) + cbar.set_label(zlabel, + rotation=90, + labelpad=0, + fontsize=_COLORBAR_LABEL_FONT_SIZE) + + if wall_file: + resolved_wall = wall_file if wall_file[0] == "/" else path + wall_file + wall_data = np.loadtxt(resolved_wall, delimiter=",") + ax.plot(wall_data[:, 0], wall_data[:, 1], color="grey") + + ax.set_xlabel(xlabel, fontsize=_XY_LABEL_FONT_SIZE) + ax.set_ylabel(ylabel, fontsize=_XY_LABEL_FONT_SIZE) + ax.set_title(title, fontsize=_TITLE_FONT_SIZE) + if xlim: + ax.set_xlim(xlim[0], xlim[1]) + if ylim: + ax.set_ylim(ylim[0], ylim[1]) + utils.set_tick_font_size(ax, _TICK_FONT_SIZE) + + if saveas: + fig.savefig(saveas) + if show: + plt.show() + + return fig diff --git a/src/postgkyl/diagnostics/gk/particle_balance.py b/src/postgkyl/diagnostics/gk/particle_balance.py new file mode 100644 index 00000000..8f4a2c3b --- /dev/null +++ b/src/postgkyl/diagnostics/gk/particle_balance.py @@ -0,0 +1,280 @@ +"""Gyrokinetic particle-balance diagnostic. + +Ported from ``src_bak/postgkyl/apps/gk_particle_balance.py``. Same shape as +:mod:`postgkyl.diagnostics.gk.energy_balance`, but for a single +species and the M0 (density) moment, with no field/apar-energy terms:: + + N_err = S - bflux - df/dt +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import matplotlib.pyplot as plt +import numpy as np + +from . import utils + +_DIRS = ("x", "y", "z") +_EDGES = ("lower", "upper") +_LINE_STYLES = ("-", "--", ":", "-.") +_XY_LABEL_FONT_SIZE = 17 +_TITLE_FONT_SIZE = 17 +_TICK_FONT_SIZE = 14 +_LEGEND_FONT_SIZE = 14 + +# Integrated-moments files store (M0, M1, M2, ...) per component; particle +# balance uses the M0 (density) moment, index 0. +_DENSITY_MOMENT = 0 + + +@dataclass(frozen=True) +class ParticleBalanceTraces: + """Computed particle-balance time traces (all 1-D, aligned to ``time``). + + Attributes: + time: Time stamps of the ``fdot`` trace. + fdot: Rate of change of the M0 moment, summed over blocks. + src: Rate of change from sources, or ``None`` if none was found. + bflux_tot: Rate of change from boundary particle fluxes, or ``None``. + mom_err: The particle-balance residual (``None`` when ``relative_error``). + mom_err_norm: The *relative* residual (only set when + ``relative_error=True``). + """ + + time: np.ndarray + fdot: np.ndarray + src: np.ndarray | None + bflux_tot: np.ndarray | None + mom_err: np.ndarray | None + mom_err_norm: np.ndarray | None = None + + +def _accumulate(target: np.ndarray | None, addend) -> np.ndarray: + """Sum ``addend`` into ``target`` (over blocks), copying on first use so + the caller's array is never mutated in place.""" + addend = np.asarray(addend) + return addend.copy() if target is None else target + addend + + +def particle_balance_error(fdot: np.ndarray, src: np.ndarray, + bflux_tot: np.ndarray) -> np.ndarray: + """The particle-balance residual: ``S - bflux - df/dt``.""" + return src - bflux_tot - fdot + + +def _block_prefix(file_prefix: str, block_idx: int) -> str: + return file_prefix.replace("*", str(block_idx)) + + +def _resolve(path: str, override: str | None, default: str, + block_idx: int) -> str: + """Resolve a file-family member's path: ``override`` (with ``*`` + substituted for the block index) if given, else the naming-convention + ``default``.""" + if override is None: + return default + return (path + override).replace("*", str(block_idx)) + + +def particle_balance( + name: str, + species: str, + *, + path: str = "./", + relative_error: bool = False, + multib: str = "-10", + fdot_file: str | None = None, + source_file: str | None = None, + bflux_files: dict[str, str] | None = None, + f_file: str | None = None, + dt_file: str | None = None, + logy: bool = False, + absy: bool = False, + xlabel: str = "Time (s)", + ylabel: str | None = None, + title: str | None = None, + indent_left: float = 0.0, + add_width: float = 0.0, + show: bool = False, + saveas: str | None = None, +) -> tuple[plt.Figure, ParticleBalanceTraces]: + """Plot (and compute) the particle balance of a single species. + + Requires ``-_fdot_integrated_moms.gkyl``; and (only if the + run had sources or non-periodic boundaries) + ``-_source_integrated_moms.gkyl`` and + ``-_bflux__integrated_HamiltonianMoments + .gkyl`` files. If ``relative_error`` is requested, + ``-_integrated_moms.gkyl`` and ``-dt.gkyl`` are also + required. + + Args: + name: Simulation name (also the file prefix). + species: Species name. + path: Directory holding the simulation output. + relative_error: Plot the relative error instead of every balance term. + multib: ``"-10"`` (default) for a single block; ``"-1"`` to discover + every block; otherwise a comma list or ``'start:stop[:step]'`` slice + of block indices. + fdot_file: Explicit distribution derivative moments path override. + source_file: Explicit source moments path override. + bflux_files: Optional per-boundary path overrides, keyed by + ``""``; unlisted boundaries use the naming + convention. + f_file: Explicit integrated distribution moments path override. + dt_file: Explicit time-step history path override. In path overrides, + ``*`` stands for the block index. + logy: Log-scale the y axis. + absy: Take the absolute value of every trace before plotting. + xlabel: Horizontal-axis label. + ylabel: Vertical-axis label. + title: Figure title. + indent_left: Horizontal axes-position adjustment in figure units. + add_width: Axes-width adjustment in figure units. + show: Call ``plt.show()`` before returning. + saveas: If given, save the figure to this path. + + Returns: + ``(figure, traces)``. + + Raises: + FileNotFoundError: if a required file family is missing. + """ + path = path.rstrip("/") + "/" + bflux_files = bflux_files or {} + + file_prefix = f"{path}{name}-" if multib == "-10" else f"{path}{name}_b*-" + probe = fdot_file or (file_prefix + species + "_fdot_integrated_moms.gkyl") + blocks = utils.get_block_indices(multib, probe) + + fig = plt.figure(figsize=(7.5, 4.5)) + ax = fig.add_axes([0.11 + indent_left, 0.15, 0.87 + add_width, 0.78]) + ax.plot([-1.0, 1.0], [0.0, 0.0], color="grey", linestyle=":", linewidth=1) + + absy_func = np.abs if absy else (lambda v: v) + + fdot = src = bflux_tot = None + has_src = has_bflux = False + time_fdot = time_bflux_tot = None + + for block_idx in blocks: + block_prefix = _block_prefix(file_prefix, block_idx) + + fdot_name = _resolve(path, fdot_file, + block_prefix + species + "_fdot_integrated_moms.gkyl", + block_idx) + found, t, v, _ = utils.read_time_trace_if_present(fdot_name) + if not found: + raise FileNotFoundError(f"Required file not found: {fdot_name}") + time_fdot = t + fdot_pb = v[:, _DENSITY_MOMENT] + + src_name = _resolve(path, source_file, + block_prefix + species + "_source_integrated_moms.gkyl", + block_idx) + has_src, t, v, _ = utils.read_time_trace_if_present(src_name) + src_pb = v[:, _DENSITY_MOMENT] if has_src else 0.0 * fdot_pb + + bflux_terms = [] + for d in _DIRS: + for e in _EDGES: + key = d + e + bf_name = _resolve( + path, bflux_files.get(key), block_prefix + species + + f"_bflux_{d}{e}_integrated_HamiltonianMoments.gkyl", block_idx) + found_b, t, v, _ = utils.read_time_trace_if_present(bf_name) + if found_b: + has_bflux = True + time_bflux_tot = t + bflux_terms.append(v[:, _DENSITY_MOMENT]) + bflux_pb = sum(bflux_terms) if bflux_terms else 0.0 * fdot_pb + + fdot = _accumulate(fdot, fdot_pb) + src = _accumulate(src, src_pb) + bflux_tot = _accumulate(bflux_tot, bflux_pb) + + legend_handles = [] + legend_strings = [] + + if not relative_error: + src = src.copy() + src[0] = 0.0 # No fdot/bflux contribution at t=0. + + mom_err = particle_balance_error(fdot, src, bflux_tot) + + if has_src: + h, = ax.plot(time_fdot, absy_func(src), linestyle=_LINE_STYLES[2]) + legend_handles.append(h) + legend_strings.append(r"$\mathcal{S}$") + if has_bflux: + h, = ax.plot(time_bflux_tot, + absy_func(-bflux_tot), + linestyle=_LINE_STYLES[1]) + legend_handles.append(h) + legend_strings.append( + r"$-\int_{\partial \Omega}\mathrm{d}\mathbf{S}\cdot\mathbf{\dot{R}}f$" + ) + h, = ax.plot(time_fdot, absy_func(-fdot), linestyle=_LINE_STYLES[0]) + legend_handles.append(h) + legend_strings.append(r"$-\dot{f}$") + h, = ax.plot(time_fdot, absy_func(mom_err), linestyle=_LINE_STYLES[3]) + legend_handles.append(h) + legend_strings.append(r"$E_{\dot{\mathcal{N}}}=$" + "".join(legend_strings)) + + ax.legend(legend_handles, + legend_strings, + fontsize=_LEGEND_FONT_SIZE, + frameon=False) + + ylabel_string = ylabel or "" + title_string = title or r"Particle balance" + mom_err_norm = None + else: + dt_name = _resolve(path, dt_file, + file_prefix.replace("_b*", "") + "dt.gkyl", 0) + _, time_dt, dt, _ = utils.read_time_trace_if_present(dt_name) + + distf = None + for block_idx in blocks: + block_prefix = _block_prefix(file_prefix, block_idx) + f_name = _resolve(path, f_file, + block_prefix + species + "_integrated_moms.gkyl", + block_idx) + _, t, v, _ = utils.read_time_trace_if_present(f_name) + distf = _accumulate(distf, v[:, _DENSITY_MOMENT]) + + fdot, src, bflux_tot, distf = fdot[1:], src[1:], bflux_tot[1:], distf[1:] + mom_err = particle_balance_error(fdot, src, bflux_tot) + mom_err_norm = mom_err * dt / distf + + ax.plot(time_dt, absy_func(mom_err_norm)) + + ylabel_string = ylabel or r"$E_{\dot{\mathcal{N}}}~\Delta t/\mathcal{N}$" + title_string = title or r"Relative error in particle conservation" + mom_err = None + + if logy: + ax.set_yscale("log") + if absy and ylabel_string: + ylabel_string = r"|" + ylabel_string + r"|" + + ax.set_xlabel(xlabel, fontsize=_XY_LABEL_FONT_SIZE) + ax.set_ylabel(ylabel_string, fontsize=_XY_LABEL_FONT_SIZE) + ax.set_title(title_string, fontsize=_TITLE_FONT_SIZE) + ax.set_xlim(time_fdot[0], time_fdot[-1]) + utils.set_tick_font_size(ax, _TICK_FONT_SIZE) + + if saveas: + fig.savefig(saveas) + if show: + plt.show() + + traces = ParticleBalanceTraces(time=time_fdot, + fdot=fdot, + src=src if has_src else None, + bflux_tot=bflux_tot if has_bflux else None, + mom_err=mom_err, + mom_err_norm=mom_err_norm) + return fig, traces diff --git a/src/postgkyl/diagnostics/gk/quantities.py b/src/postgkyl/diagnostics/gk/quantities.py new file mode 100644 index 00000000..18cbaeb9 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/quantities.py @@ -0,0 +1,609 @@ +"""Gyrokinetic derived-quantity physics -- the ``fetch_*`` functions behind the +quantity registry. + +Ported from ``src_bak/postgkyl/gk/gk_quantities/fetch_funcs.py``. Every +``fetch_*`` there computed through ``GkeyllDGops`` -- a ``ctypes`` binding +that is dead in this tree (rule #2). Rewired here onto the new surface: +every fetch function **interpolates its inputs first** +(:meth:`~postgkyl.gdata.gdata.GData.interpolate`, the sanctioned "evaluation" +bridge -- REFACTOR_GKEYLL_FFI.md's field domain) and then computes with +plain NumPy on the interpolated values, exactly like every sibling equation +module (``five_moment``, ``ten_moment``, ``mhd``, ...). This is a deliberate +divergence from a literal "stay modal and call the weak kernels" port: +extracting one physical field's coefficients out of a *packed* multi-field +source file (``M0M1M2``, ``BiMaxwellianMoments``, ``HamiltonianMoments``, ...) +has no primitive reachable from this layer's allowed imports (``gdatastate``, +``operations``, ``numerics``, ``api`` -- not ``dg``/``gpython``; only ``operations.select`` +could slice a component, and it unconditionally refuses gkyl-backed data). +Interpolating first sidesteps that gap entirely and matches the one +established working pattern in this codebase; see the layer-12 report for +the full trade-off discussion. Physical constants come from +``scipy.constants`` (rule #13), not a re-typed ``gk/gkeyll_const.py`` table. + +Naming keys (matching ``src_bak`` so the registry mapping in ``registry.py`` +stays recognizable): + s#: source #, c#: component #, add/sub/mul/div: the combining operator, + pos/neg: the plus/minus term of a curvilinear cross product. +""" + +from __future__ import annotations + +import operator +from typing import TYPE_CHECKING + +import numpy as np +from scipy import constants + +from postgkyl import operations + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _get_ctx_val(gdata: "GDataState", key: str, **kwargs): + """A value (or one value per species) for ``key``: ``kwargs[key]`` + (an explicit ``--extra`` override) wins over ``gdata.ctx[key]`` (the + file's own attribute), which wins over raising. + + ``kwargs[key]`` may be a single value (applies to every species) or a + list/tuple with one entry per species, indexed by ``kwargs + ['species_idx']`` -- the position :meth:`~postgkyl.diagnostics. + gk.quantity.GkQuantity.fetch_multi`/``load_quantity`` stamp + onto ``extra`` for the species currently being resolved. + """ + if key in kwargs: + val = kwargs[key] + if not isinstance(val, (list, tuple)): + return val + species_idx = kwargs.get("species_idx") + if species_idx is None: + raise KeyError( + f"fetch function: '--extra {key}=' was given {len(val)} values " + "but this quantity is not resolved per species here, so there is " + "no way to tell which one to use. Pass a single value instead.") + if species_idx >= len(val): + species = kwargs.get("species") + raise ValueError( + f"fetch function: '--extra {key}=' was given only {len(val)} " + f"values but species #{species_idx}" + f"{f' ({species})' if species else ''} was requested. Give one " + "value per species, in the order of '--species'.") + return val[species_idx] + if gdata.ctx.get(key) is not None: + return gdata.ctx[key] + raise KeyError( + f"fetch function: context key '{key}' not found in the dataset; " + f"pass it as '--extra {key}=', or as one value per species " + f"with '--extra {key}=,,...'.") + + +def _ensure_interpolated(d: "GDataState") -> "GDataState": + """Interpolate ``d`` onto the field domain unless it already is. + + Uses the ``operations.interpolate`` verb directly (rather than the fluent + ``GData.interpolate()``) so this works on any ``GDataState``, not just the + fluent subclass -- these functions receive whatever + ``GkQuantity.get_src_gdata`` hands them. + """ + if d.ctx.get("interpolated"): + return d + return operations.interpolate(d) + + +def _component(d: "GDataState", comp: int | None) -> "GDataState": + """Interpolate ``d`` and select physical component ``comp`` (all if None).""" + interpolated = _ensure_interpolated(d) + return interpolated if comp is None else operations.select(interpolated, + comp=comp) + + +# --------------------------------------------------- generic fetch factories +def _make_fetch_comp(icomp: int | None): + """A fetch function that extracts the ``icomp``-th physical component.""" + + def fetch(gdatas, **kwargs): + return _component(gdatas[0], icomp) + + fetch.__name__ = f"fetch_comp{icomp}" if icomp is not None else "fetch_compAll" + return fetch + + +def _make_fetch_binop(si: int, ci: int, sj: int, cj: int, op): + """A fetch function combining component ``ci`` of source ``si`` with + component ``cj`` of source ``sj`` via ``op`` (both interpolated first).""" + + def fetch(gdatas, **kwargs): + a = _component(gdatas[si], ci) + b = _component(gdatas[sj], cj) + return a._result(a.grid, op(a.values, b.values)) + + fetch.__name__ = f"fetch_s{si}c{ci}_{op.__name__}_s{sj}c{cj}" + return fetch + + +# Extract a single component. +fetch_s0cAll = _make_fetch_comp(None) +fetch_s0c0 = _make_fetch_comp(0) +fetch_s0c1 = _make_fetch_comp(1) +fetch_s0c2 = _make_fetch_comp(2) +fetch_s0c3 = _make_fetch_comp(3) + +# Combine components across (possibly different) sources. +fetch_s0c0_add_s1c0 = _make_fetch_binop(0, 0, 1, 0, operator.add) +fetch_s0c2_add_s0c3 = _make_fetch_binop(0, 2, 0, 3, operator.add) +fetch_s0c0_sub_s1c0 = _make_fetch_binop(0, 0, 1, 0, operator.sub) +fetch_s0c0_mul_s1c0 = _make_fetch_binop(0, 0, 1, 0, operator.mul) +fetch_s0c0_mul_s0c1 = _make_fetch_binop(0, 0, 0, 1, operator.mul) +fetch_s1c0_div_s0c0 = _make_fetch_binop(1, 0, 0, 0, operator.truediv) + + +# ------------------------------------------------------------------ moments +def fetch_M1_from_H(gdatas, **kwargs): + """M1 from the Hamiltonian moments: ``mass**-1 * (comp0 * comp1)``.""" + hmom = _ensure_interpolated(gdatas[0]) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + values = hmom.values[..., 0, np.newaxis] * hmom.values[..., 1, np.newaxis] + return hmom._result(hmom.grid, values / mass) + + +def fetch_Tpar_from_BiMax(gdatas, **kwargs): + """Tpar from BiMaxwellian moments: ``mass * comp2``.""" + Tpar = fetch_s0c2(gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + return Tpar._result(Tpar.grid, mass * Tpar.values) + + +def fetch_Tpar_from_M0_M1_M2par(gdatas, **kwargs): + """``upar*M1 + M0*Tpar/m = M2par`` => ``Tpar = m*(M2par - upar*M1)/M0``.""" + m0, m1, m2par = (_ensure_interpolated(g) for g in gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + upar = m1.values / m0.values + values = mass * (m2par.values - upar * m1.values) / m0.values + return m0._result(m0.grid, values) + + +def fetch_Tperp_from_BiMax(gdatas, **kwargs): + """Tperp from BiMaxwellian moments: ``mass * comp3``.""" + Tperp = fetch_s0c3(gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + return Tperp._result(Tperp.grid, mass * Tperp.values) + + +def fetch_Tperp_from_M0_M2perp(gdatas, **kwargs): + """``Tperp = 0.5 * mass * (M2perp / M0)``.""" + Tperp = fetch_s1c0_div_s0c0(gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + return Tperp._result(Tperp.grid, 0.5 * mass * Tperp.values) + + +def fetch_temp_from_Max(gdatas, **kwargs): + """temp from Maxwellian moments: ``mass * comp2``.""" + temp = fetch_s0c2(gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + return temp._result(temp.grid, mass * temp.values) + + +def fetch_temp_from_Tpar_Tperp(gdatas, **kwargs): + """``temp = (Tpar + 2*Tperp) / 3``.""" + Tpar, Tperp = (_ensure_interpolated(g) for g in gdatas) + values = (Tpar.values + 2.0 * Tperp.values) / 3.0 + return Tpar._result(Tpar.grid, values) + + +def fetch_press_from_Max(gdatas, **kwargs): + """Pressure from Maxwellian moments: ``press = mass * comp0 * comp2``.""" + maxmom = _ensure_interpolated(gdatas[0]) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + values = mass * maxmom.values[..., 0, np.newaxis] * maxmom.values[..., 2, + np.newaxis] + return maxmom._result(maxmom.grid, values) + + +def fetch_press_from_BiMax(gdatas, **kwargs): + """Pressure from BiMaxwellian moments: ``press = comp0 * mass*(Tpar+2Tperp)/3``.""" + bimax = _ensure_interpolated(gdatas[0]) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + Tpar_vals = bimax.values[..., 2, np.newaxis] + Tperp_vals = bimax.values[..., 3, np.newaxis] + temp_vals = mass * (Tpar_vals + 2.0 * Tperp_vals) / 3.0 + values = bimax.values[..., 0, np.newaxis] * temp_vals + return bimax._result(bimax.grid, values) + + +def fetch_press_p(gdatas, **kwargs): + """Perpendicular/parallel pressure in J/m^3: ``p_p = n * T_p``.""" + m0, Tp = (_ensure_interpolated(g) for g in gdatas) + return m0._result(m0.grid, m0.values * Tp.values) + + +def _make_fetch_q(name: str): + """Return a fetch function for the lab-frame parallel flux of the + parallel (``name='par'``) or perpendicular (``name='perp'``) kinetic + energy:: + + q_par = (m/2)*M3par = (m/2) int(vpar^3 f) dv, + q_perp = (m/2)*M3perp = (m/2) int(vpar*vperp^2 f) dv, + + so that ``q_par + q_perp`` is the parallel flux of the total kinetic + energy. Both are in W/m^2 (kg/s^3). ``gdatas``: ``[M3par]`` or + ``[M3perp]``. + """ + + def fetch(gdatas, **kwargs): + m3 = _ensure_interpolated(gdatas[0]) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + return m3._result(m3.grid, 0.5 * mass * m3.values) + + fetch.__name__ = f"fetch_q{name}" + return fetch + + +fetch_qpar = _make_fetch_q("par") +fetch_qperp = _make_fetch_q("perp") + + +def _make_fetch_q_fluid(name: str): + """Return a fetch function for the parallel/perpendicular heat flux in + the fluid (drift) frame -- the energy carried by the random part of the + motion, ``u = M1/M0`` being the parallel drift speed:: + + q_par = (m/2) int (vpar-u)^3 f dv + = (m/2) [M3par - 3*u*M2par + 3*u^2*M1 - u^3*M0] + = (m/2) [M3par - 3*u*M2par + 2*u^2*M1], + q_perp = (m/2) int (vpar-u)*vperp^2 f dv + = (m/2) [M3perp - u*M2perp]. + + ``gdatas`` (in this order): ``[M0, M1, M2par, M3par]`` or + ``[M0, M1, M2perp, M3perp]``. + """ + is_par = name == "par" + + def fetch(gdatas, **kwargs): + m0, m1, m2, m3 = (_ensure_interpolated(g) for g in gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + + upar = m1.values / m0.values + u_m2 = upar * m2.values + + if is_par: + values = m3.values - 3.0 * u_m2 + 2.0 * upar**2 * m1.values + else: + values = m3.values - u_m2 + + return m0._result(m0.grid, 0.5 * mass * values) + + fetch.__name__ = f"fetch_q{name}_fluid" + return fetch + + +fetch_qpar_fluid = _make_fetch_q_fluid("par") +fetch_qperp_fluid = _make_fetch_q_fluid("perp") + + +def fetch_vt(gdatas, **kwargs): + """Thermal speed ``vt = sqrt(T/m)`` (m/s), ``m`` the requested species' + mass. ``gdatas``: ``[temp]`` (temperature, in Joules).""" + temp = _ensure_interpolated(gdatas[0]) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + return temp._result(temp.grid, np.sqrt(temp.values / mass)) + + +def fetch_larmor_radius(gdatas, **kwargs): + """Species Larmor (gyro-)radius: ``rho = sqrt(m*T)/(|q|*B)``. ``gdatas``: + ``[temp, bmag]``.""" + temp, bmag = (_ensure_interpolated(g) for g in gdatas) + mass = _get_ctx_val(gdatas[0], "mass", **kwargs) + charge = abs(_get_ctx_val(gdatas[0], "charge", **kwargs)) + values = np.sqrt(mass * temp.values) / (charge * bmag.values) + return temp._result(temp.grid, values) + + +def fetch_debye_length(gdatas, **kwargs): + """Species-wise Debye length: ``lambda_D = sqrt(eps0*T/(n*q^2))``. + ``gdatas``: ``[temp, M0]``.""" + temp, m0 = (_ensure_interpolated(g) for g in gdatas) + charge = _get_ctx_val(gdatas[0], "charge", **kwargs) + values = np.sqrt(constants.epsilon_0 * temp.values / (m0.values * charge**2)) + return temp._result(temp.grid, values) + + +def _split_elc_ions(gdatas, quantity: str, **kwargs): + """Split the per-species sources of a multi-species quantity into the + electron entry and the ion entries, by the sign of each species' charge. + + ``gdatas[i]`` is species ``i``'s resolved source list (as + :meth:`~postgkyl.diagnostics.gk.quantity.GkQuantity.fetch_multi` + hands it to an ``is_multi_species`` fetch function); each entry's + ``mass``/``charge`` is resolved with ``species_idx=i`` so a per-species + ``--extra`` array picks the right one. + """ + species_names = kwargs.get("species", []) + if len(species_names) != len(gdatas): + species_names = [f"#{i}" for i in range(len(gdatas))] + + elcs, ions = [], [] + for species_idx, (name, srcs) in enumerate(zip(species_names, gdatas)): + species_kwargs = dict(kwargs, species_idx=species_idx, species=name) + entry = { + "name": name, + "srcs": [_ensure_interpolated(s) for s in srcs], + "mass": _get_ctx_val(srcs[0], "mass", **species_kwargs), + "charge": _get_ctx_val(srcs[0], "charge", **species_kwargs), + } + (elcs if entry["charge"] < 0.0 else ions).append(entry) + + if len(elcs) != 1: + raise ValueError( + f"{quantity}: expected exactly one negatively charged (electron) " + f"species but found {len(elcs)} in {list(species_names)}.") + if not ions: + raise ValueError( + f"{quantity}: found no positively charged (ion) species in " + f"{list(species_names)}.") + return elcs[0], ions + + +def _weighted_sum(entries, weights, comp: int) -> "GDataState": + """Sum the ``comp``-th (already-interpolated) source of each species, + each scaled by a scalar weight.""" + base = entries[0]["srcs"][comp] + total = sum(w * e["srcs"][comp].values for e, w in zip(entries, weights)) + return base._result(base.grid, total) + + +def _fetch_c_s_ion_acoustic(gdatas, **kwargs): + """Ion-acoustic sound speed (wave perspective), for the Bohm criterion + and sheath/presheath matching:: + + c_s = sqrt( T_e * sum_j(n_j*Z_j^2/m_j) / sum_j(n_j*Z_j) ) + + summing over the ion species ``j``, with ``Z_j = q_j/e`` the ion charge + state. + """ + elc, ions = _split_elc_ions(gdatas, "fetch_c_s(kind=ion_acoustic)", **kwargs) + + e = constants.elementary_charge + charge_states = [ion["charge"] / e for ion in ions] + + numer = _weighted_sum( + ions, [z**2 / ion["mass"] for z, ion in zip(charge_states, ions)], 0) + denom = _weighted_sum(ions, charge_states, 0) + + temp_e = elc["srcs"][1] + values = np.sqrt(temp_e.values * numer.values / denom.values) + return temp_e._result(temp_e.grid, values) + + +def _fetch_c_s_thermo(gdatas, **kwargs): + """Thermodynamic sound speed (bulk fluid perspective), for Mach numbers + and acoustic propagation in the core/SOL:: + + c_s = sqrt( (gamma_e*n_e*T_e + sum_j(gamma_j*n_j*T_j)) / sum_j(n_j*m_j) ) + + summing over the ion species ``j``. Default ``gamma_e=1``, ``gamma_i=3``, + overridable via ``--extra``. + """ + elc, ions = _split_elc_ions(gdatas, "fetch_c_s(kind=thermo)", **kwargs) + + gamma_e = float(kwargs.get("gamma_e", 1.0)) + gamma_i = float(kwargs.get("gamma_i", 3.0)) + + m0_e, temp_e = elc["srcs"] + numer_vals = gamma_e * m0_e.values * temp_e.values + for ion in ions: + m0_i, temp_i = ion["srcs"] + numer_vals = numer_vals + gamma_i * m0_i.values * temp_i.values + + denom = _weighted_sum(ions, [ion["mass"] for ion in ions], 0) + values = np.sqrt(numer_vals / denom.values) + return temp_e._result(temp_e.grid, values) + + +def fetch_c_s(gdatas, **kwargs): + """Sound speed (m/s), combining the electrons and every ion species. + ``gdatas`` has one ``[M0, temp]`` source list per species, in the order + requested, e.g. ``pgkyl gk_load_quantity --quantity c_s + --species elc,ion1,ion2 ...``. + Electrons and ions are told apart by the sign of each species' charge + attribute, so the species may be named anything. + + Two definitions are available through ``--extra kind=``: + ``ion_acoustic``: the wave/Bohm-criterion sound speed, + ``c_s = sqrt(T_e*sum_j(n_j*Z_j^2/m_j)/sum_j(n_j*Z_j))``. + ``thermo`` (default): the bulk-fluid sound speed, + ``c_s = sqrt((gamma_e*n_e*T_e + sum_j(gamma_j*n_j*T_j))/sum_j(n_j*m_j))``, + with ``gamma_e``/``gamma_i`` settable via ``--extra`` (default 1, 3). + """ + c_s_kinds = { + "ion_acoustic": _fetch_c_s_ion_acoustic, + "thermo": _fetch_c_s_thermo, + } + kind = str(kwargs.get("kind", "thermo")) + if kind not in c_s_kinds: + raise ValueError( + f"fetch_c_s: unknown kind '{kind}'. Select one with '--extra " + f"kind=' from: {', '.join(sorted(c_s_kinds))}.") + return c_s_kinds[kind](gdatas, **kwargs) + + +def fetch_beta_from_bmag_press(gdatas, **kwargs): + """``beta = 2*mu_0*press / bmag**2``.""" + bmag, press = (_ensure_interpolated(g) for g in gdatas) + values = 2.0 * constants.mu_0 * press.values / bmag.values**2 + return bmag._result(bmag.grid, values) + + +# ------------------------------------------------------------ drift speeds +def _b_cross_grad_div_b_component(scalar: "GDataState", + jacobtot_inv: "GDataState", b_i: "GDataState", + comp: int) -> "GDataState": + """The ``comp``-th component of ``b x grad(f) / (J B)``. + + ``(b x grad f)_k / B = epsilon_{ijk} * b_i * d(f)/dx^j / (J B)``, where + ``epsilon_{ijk}`` is the Levi-Civita tensor, ``f`` a scalar field, ``b_i`` + the covariant components of a vector field. The gradient is the numerical + (post-``interpolate()``) one (``operations.differentiate``); see + ``differentiate-decision.md`` -- an exact modal derivative needs a shim + addition out of scope for this layer. + + Args: + scalar: Scalar field ``f`` to differentiate; interpolated internally. + jacobtot_inv: Inverse of the total-coordinate-transformation Jacobian. + b_i: Covariant components of the unit vector field ``b``. + comp: 0-based component ``k`` of the cross product (``< 3``). + + Raises: + KeyError: if ``comp`` is not 0, 1, or 2. + """ + f = _ensure_interpolated(scalar) + cdim = f.num_dims + + diff_dir_pos = bi_c_pos = 0 + diff_dir_neg = bi_c_neg = 0 + calc_term = [True, True] + if comp == 0: + diff_dir_neg = bi_c_pos = 1 + diff_dir_pos = bi_c_neg = cdim - 1 + if cdim < 3: + calc_term = [True, False] + elif comp == 1: + bi_c_pos, bi_c_neg = 2, 0 + diff_dir_neg, diff_dir_pos = cdim - 1, 0 + if cdim == 1: + calc_term = [False, True] + elif comp == 2: + diff_dir_neg = bi_c_pos = 0 + diff_dir_pos = bi_c_neg = 1 + if cdim == 1: + calc_term = [False, False] + elif cdim == 2: + calc_term = [False, True] + else: + raise KeyError("comp must be 0, 1, or 2.") + + b_i_i = _ensure_interpolated(b_i) + jacobtot_inv_i = _ensure_interpolated(jacobtot_inv) + + pos_term = np.zeros_like(f.values) + neg_term = np.zeros_like(f.values) + if calc_term[0]: + d_pos = operations.differentiate(f, direction=diff_dir_pos) + pos_term = d_pos.values * b_i_i.values[..., bi_c_pos, np.newaxis] + if calc_term[1]: + d_neg = operations.differentiate(f, direction=diff_dir_neg) + neg_term = -d_neg.values * b_i_i.values[..., bi_c_neg, np.newaxis] + + values = (pos_term + neg_term) * jacobtot_inv_i.values + return f._result(f.grid, values) + + +def fetch_ExB_vel(gdatas, **kwargs): + """``v_{E,k} = epsilon_{ijk}/(J B) * b_i * d(phi)/dx^j`` (``dir`` selects k). + + ``gdatas``: ``(jacobtot_inv, bmag, b_i, phi)``. + """ + if "dir" not in kwargs: + raise KeyError("fetch_ExB_vel: select the k-th component with dir=.") + jacobtot_inv, _bmag, b_i, phi = gdatas + return _b_cross_grad_div_b_component(phi, jacobtot_inv, b_i, kwargs["dir"]) + + +def fetch_gradB_vel(gdatas, **kwargs): + """``v_gradB,k = Tperp/(q B) * epsilon_{ijk} * b_i * d(B)/dx^j / (J B)``. + + ``gdatas``: ``(jacobtot_inv, bmag, b_i, Tperp)``. + """ + if "dir" not in kwargs: + raise KeyError("fetch_gradB_vel: select the k-th component with dir=.") + jacobtot_inv, bmag, b_i, Tperp = gdatas + out = _b_cross_grad_div_b_component(bmag, jacobtot_inv, b_i, kwargs["dir"]) + bmag_i = _ensure_interpolated(bmag) + Tperp_i = _ensure_interpolated(Tperp) + charge = _get_ctx_val(Tperp, "charge", **kwargs) + values = out.values * Tperp_i.values / bmag_i.values / charge + return out._result(out.grid, values) + + +def fetch_diamag_vel(gdatas, **kwargs): + """``v_diamag,k = 1/(q n) epsilon_{ijk} b_i * d(pperp)/dx^j / (J B)``. + + ``gdatas``: ``(jacobtot_inv, bmag, b_i, m0, pressperp)``. + """ + if "dir" not in kwargs: + raise KeyError( + "fetch_diamag_vel: select the k-th component with dir=.") + jacobtot_inv, bmag, b_i, m0, pressperp = gdatas + out = _b_cross_grad_div_b_component(pressperp, jacobtot_inv, b_i, + kwargs["dir"]) + m0_i = _ensure_interpolated(m0) + charge = _get_ctx_val(pressperp, "charge", **kwargs) + values = out.values / m0_i.values / charge + return out._result(out.grid, values) + + +# --------------------------------------------------------- phase space (f) +def load_distf(gdatas, **kwargs): + """Loader for the registry ``distf`` quantity: wraps + :func:`~postgkyl.diagnostics.gk.distf.load_distf` with + defaults tailored to registry use (never interpolate further, convert + velocity coordinates by default). Extra keyword overrides (via + ``**extra`` on :func:`~postgkyl.diagnostics.gk.load_quantity. + load_quantity`): ``suffix``, ``c2p_vel``, ``mc2nu``, ``mapc2p``, + ``block``. + """ + from .distf import load_distf + from .utils import dict_get_bool + + prefix = kwargs.get("path", "").rstrip("/") + "/" + kwargs.get("name", "") + extra = { + k: v + for k, v in kwargs.items() + if k not in ("path", "name", "species", "frame") + } + + return load_distf( + name=prefix, + species=kwargs.get("species", ""), + frame=int(kwargs.get("frame", 0)), + suffix=str(extra.get("suffix", "")), + use_c2p_vel=dict_get_bool(extra, "c2p_vel", True), + use_mc2nu=dict_get_bool(extra, "mc2nu", False), + use_mapc2p=dict_get_bool(extra, "mapc2p", False), + block_idx=extra.get("block", None), + num_interp=0, + ) + + +# ----------------------------------------------------- normalized quantities +def _make_fetch_q_norm(name: str): + """Return a fetch function for a heat flux normalized by the + free-streaming estimate ``n*T*c_s``: ``q_norm = q / (n*T*c_s)``. + ``gdatas`` (in this order): ``[q, M0, temp, c_s]``. + """ + + def fetch(gdatas, **kwargs): + q, m0, temp, c_s = (_ensure_interpolated(g) for g in gdatas) + values = q.values / (m0.values * temp.values * c_s.values) + return q._result(q.grid, values) + + fetch.__name__ = f"fetch_q{name}_norm" + return fetch + + +fetch_qpar_norm = _make_fetch_q_norm("par") +fetch_qperp_norm = _make_fetch_q_norm("perp") + + +def fetch_rho_over_lambda(gdatas, **kwargs): + """Ratio of the species Larmor radius to its Debye length: + ``rho/lambda_D``. ``gdatas``: ``[rho, lambda_D]``.""" + rho, lambda_d = (_ensure_interpolated(g) for g in gdatas) + return rho._result(rho.grid, rho.values / lambda_d.values) + + +def fetch_phi_norm(gdatas, **kwargs): + """Normalized electrostatic potential: ``phi_norm = e*phi/T_e``. + ``gdatas``: ``[phi, temp]``.""" + phi, temp = (_ensure_interpolated(g) for g in gdatas) + values = constants.elementary_charge * phi.values / temp.values + return phi._result(phi.grid, values) diff --git a/src/postgkyl/diagnostics/gk/quantity.py b/src/postgkyl/diagnostics/gk/quantity.py new file mode 100644 index 00000000..172f33b7 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/quantity.py @@ -0,0 +1,314 @@ +"""``GkQuantity`` -- a registered gyrokinetic quantity, and its registry. + +Ported from ``src_bak/postgkyl/gk/gk_quantities/gkquantity.py``. A quantity +names one or more *source combinations* (files and/or other, already- +registered ``GkQuantity`` objects) together with the ``fetch_func`` that +turns a resolved combination into the quantity's data. Source-combination +frame discovery calls :mod:`postgkyl.diagnostics.discovery` -- the one home +for "what outputs does this directory hold" -- instead of globbing on its +own. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Callable, TYPE_CHECKING + +from postgkyl.gdata import GData + +from .. import discovery + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +@dataclass(frozen=True) +class GkQuantity: + """A gyrokinetic quantity: one or more source combinations + fetch logic. + + Attributes: + name: Name of the quantity (the registry key). + source: List of source combinations to try, in preference order; each + combination is a list of either file-naming-convention source strings + (e.g. ``"M0"``) or nested ``GkQuantity`` (computed on demand). + fetch_func: The fetch function for each entry in ``source`` (same + index), taking the resolved list of source ``GDataState`` and + returning the quantity's ``GDataState``. + label: LaTeX-format label for plotting (``%s`` for species name or + direction). + is_time_dep: Whether the quantity is time-dependent (written in frames). + is_species_dep: Whether the quantity is species-dependent. + is_vector: Whether the quantity is a vector (multiple components, + selected via the ``dir`` extra). + is_tensor: Whether the quantity is a tensor. + is_integrated: Whether the quantity is a grid integral. + is_geo: Whether the quantity is a (frame-independent) geometry + quantity, named ``-.gkyl`` with no frame number. + is_multi_species: Whether the quantity combines several species into a + single dataset (e.g. the sound speed, which mixes the electrons and + every ion). Such a quantity is fetched once for the whole species + list rather than once per species (:meth:`get_avail_source_multi`/ + :meth:`fetch_multi`), and its fetch function receives one list of + sources per species instead of a flat list. + """ + + name: str + source: list + fetch_func: list[Callable] + label: str + is_time_dep: bool = False + is_species_dep: bool = False + is_vector: bool = False + is_tensor: bool = False + is_integrated: bool = False + is_geo: bool = False + is_multi_species: bool = False + + # ------------------------------------------------------------ internal + def _src_stem(self, path: str, name: str, species: str, src: str) -> str: + """Stem of a string source's file name, up to (not including) the frame + number (geo files have no frame, so no trailing separator).""" + if self.is_geo: + return os.path.join(path, f"{name}-{src}") + if self.is_species_dep: + src_ = f"{src}_" if src else "" + return os.path.join(path, f"{name}-{species}_{src_}") + return os.path.join(path, f"{name}-{src}_") + + def _src_file_name(self, path: str, name: str, species: str, src: str, + frame: int | None) -> str: + """Full file name for a string source at the given frame.""" + stem = self._src_stem(path, name, species, src) + if self.is_geo: + return f"{stem}.gkyl" + return f"{stem}{frame}.gkyl" + + def _avail_frames_src(self, + path: str, + name: str, + species: str, + src: str, + frames: list[int] | None = None) -> set[int]: + """Available frames for a string source's ``.gkyl`` family.""" + stem = self._src_stem(path, name, species, src) + return discovery.available_frames(stem, frames=frames) + + def _avail_combo_frames(self, + path: str, + name: str, + species: str, + frames: list[int] | None = None + ) -> tuple[int, set[int]]: + """Find the first source combination whose files all exist and share the + same set of available frames. + + Returns: + ``(combo_idx, frames_avail)``; a combination made up only of geo files + is flagged with ``frames_avail == {-1}``. + """ + frames_avail: set[int] = set() + combo_idx = 0 + for cidx, combo in enumerate(self.source): + for src in combo: + if isinstance(src, str) and self.is_geo: + if not os.path.isfile(os.path.join(path, f"{name}-{src}.gkyl")): + frames_avail = set() + break + continue + + if isinstance(src, str): + frames_avail_q = self._avail_frames_src(path, name, species, src, + frames) + else: + _, frames_avail_q = src._avail_combo_frames(path, name, species, + frames) + + if frames_avail_q == {-1}: + combo_idx = cidx + continue + + if frames_avail_q: + if not frames_avail: + frames_avail = set(frames_avail_q) + elif frames_avail_q != frames_avail: + frames_avail = set() + break + combo_idx = cidx + else: + break + else: + if not frames_avail: + frames_avail = {-1} + combo_idx = cidx + + if frames_avail: + break + return combo_idx, frames_avail + + # -------------------------------------------------------------- public + def get_label(self, + species: str | None = None, + direction: str | None = None) -> str: + """Get the display label, substituting ``%s`` with species or direction.""" + if self.is_vector: + return self.label % str( + direction) if direction is not None else self.label % "i" + if self.is_species_dep: + return self.label % str( + species[0]) if species is not None else self.label % "s" + return self.label + + def get_avail_source(self, path: str, name: str, species: str, + frame_inp: str | None) -> tuple[int, list]: + """Identify the source combination and frame list needed for this + quantity. + + Args: + path: Directory containing the simulation files. + name: Simulation name prefix. + species: Species name. + frame_inp: A single frame, a comma-separated list, or a + ``'start:stop[:step]'`` range (``None``/``':'`` means every + available frame). + + Returns: + ``(combo_idx, frames)``. + + Raises: + FileNotFoundError: if no source combination's files are found. + """ + frame_list: list[int] = [] + if frame_inp is not None: + frame_inp = frame_inp.strip() + if "," in frame_inp: + frame_list = [int(f.strip()) for f in frame_inp.split(",")] + elif ":" not in frame_inp: + frame_list = [int(frame_inp)] + + combo_idx, frames_avail = self._avail_combo_frames(path, name, species, + frame_list) + + if not frames_avail: + raise FileNotFoundError( + f"No files found for the requested quantity (path={path!r}, " + f"name={name!r}).") + + if frames_avail == {-1}: + return combo_idx, [None] + + if len(frame_list) == 0: + frames_avail_sorted = sorted(frames_avail) + parts = frame_inp.split(":") if frame_inp else [""] + lower = int(parts[0]) if parts[0] else frames_avail_sorted[0] + upper = (int(parts[1]) + if len(parts) > 1 and parts[1] else frames_avail_sorted[-1] + 1) + step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 + frame_list = [ + f for f in frames_avail_sorted + if lower <= f < upper and (f - lower) % step == 0 + ] + + return combo_idx, frame_list + + def get_src_gdata(self, src: "str | GkQuantity", path: str, name: str, + species: str, frame: int | None, **extra) -> "GDataState": + """The populated dataset for one source: a loaded file, or a nested + quantity computed from its own sources.""" + if isinstance(src, str): + return GData(self._src_file_name(path, name, species, src, frame)) + combo_idx, _ = src.get_avail_source( + path, name, species, + str(frame) if frame is not None else None) + combo = src.source[combo_idx] + fetch_func = src.fetch_func[combo_idx] + gdatas = [ + src.get_src_gdata(s, path, name, species, frame, **extra) for s in combo + ] + return fetch_func(gdatas, **extra) + + def fetch(self, path: str, name: str, species: str, frame: int | None, + combo_idx: int, **extra) -> "GDataState": + """Fetch the source files for ``combo_idx`` and compute the quantity.""" + combo = self.source[combo_idx] + fetch_func = self.fetch_func[combo_idx] + gdatas = [ + self.get_src_gdata(src, path, name, species, frame, **extra) + for src in combo + ] + extra = dict(extra, path=path, name=name, species=species, frame=frame) + return fetch_func(gdatas, **extra) + + def get_avail_source_multi(self, path: str, name: str, + species_list: list[str], + frame_inp: str | None) -> tuple[int, list]: + """Multi-species counterpart of :meth:`get_avail_source`: resolve the + source combination and frames for every species in ``species_list``, + keeping only the frames available for all of them (the quantity folds + every species into one dataset, so a frame missing for any one of them + can't be computed at all). + + Raises: + FileNotFoundError: if no frame is available for every species. + """ + combo_idx = 0 + frames_common: set[int] | None = None + for species in species_list: + combo_idx, frames = self.get_avail_source(path, name, species, frame_inp) + frames_common = (set(frames) if frames_common is None else frames_common + & set(frames)) + if not frames_common: + raise FileNotFoundError( + f"No frames are available for all of the requested species " + f"{species_list} (path={path!r}, name={name!r}).") + return combo_idx, sorted(frames_common) + + def fetch_multi(self, path: str, name: str, species_list: list[str], + frame: int | None, combo_idx: int, **extra) -> "GDataState": + """Multi-species counterpart of :meth:`fetch`, for + ``is_multi_species`` quantities. + + The fetch function is handed one list of sources per species, in the + order of ``species_list``: ``gdatas[i][j]`` is the ``j``-th source of + the ``i``-th species. Each species' sources are resolved with + ``extra['species_idx']`` set to that species' position, so a + per-species ``--extra`` array (e.g. ``mass=1,2,3``) picks the right + entry inside the sources too, not just at the top level. Species names + are passed along as ``extra['species']``. + """ + combo = self.source[combo_idx] + fetch_func = self.fetch_func[combo_idx] + gdatas = [[ + self.get_src_gdata(src, path, name, species, frame, + **dict(extra, species_idx=species_idx)) + for src in combo + ] for species_idx, species in enumerate(species_list)] + extra = dict(extra, + path=path, + name=name, + species=list(species_list), + frame=frame) + return fetch_func(gdatas, **extra) + + +class GkQuantityRegistry: + """Registry of pre-named gyrokinetic quantities.""" + + def __init__(self): + self._registry: dict[str, GkQuantity] = {} + + def register(self, quantity: GkQuantity) -> None: + """Register a new gyrokinetic quantity.""" + self._registry[quantity.name] = quantity + + def get(self, name: str) -> GkQuantity | None: + """Get a registered quantity by name, or ``None`` if unregistered.""" + return self._registry.get(name) + + def list(self) -> list[str]: + """Sorted list of all registered quantity names.""" + return sorted(self._registry) + + def has(self, name: str) -> bool: + """Whether ``name`` is registered.""" + return name in self._registry diff --git a/src/postgkyl/diagnostics/gk/registry.py b/src/postgkyl/diagnostics/gk/registry.py new file mode 100644 index 00000000..503f672e --- /dev/null +++ b/src/postgkyl/diagnostics/gk/registry.py @@ -0,0 +1,377 @@ +"""The gyrokinetic quantity registry -- populated from ``quantities.py``. + +Ported from ``src_bak/postgkyl/gk/gk_quantities/registry.py``. Each entry +names its preferred source combinations (in order) and the fetch function +for each; :func:`~postgkyl.diagnostics.gk.quantity.GkQuantity. +get_avail_source` picks the first combination whose files are actually +present on disk. +""" + +from __future__ import annotations + +from . import quantities as ff +from .quantity import GkQuantity, GkQuantityRegistry + +gk_quant_registry = GkQuantityRegistry() + +# ----------------------------------------- scalar geometric quantities (geo) +_geo_int_jacobgeo = GkQuantity(name="geo_int_jacobgeo", + source=[["geo_int_jacobgeo"]], + fetch_func=[ff.fetch_s0c0], + label=r"$J$", + is_geo=True) +gk_quant_registry.register(_geo_int_jacobgeo) + +_geo_int_jacobgeo_inv = GkQuantity(name="geo_int_jacobgeo_inv", + source=[["geo_int_jacobgeo_inv"]], + fetch_func=[ff.fetch_s0c0], + label=r"$J^{-1}$", + is_geo=True) +gk_quant_registry.register(_geo_int_jacobgeo_inv) + +_geo_int_jacobtot = GkQuantity(name="geo_int_jacobtot", + source=[["geo_int_jacobtot"]], + fetch_func=[ff.fetch_s0c0], + label=r"$J$", + is_geo=True) +gk_quant_registry.register(_geo_int_jacobtot) + +_geo_int_jacobtot_inv = GkQuantity(name="geo_int_jacobtot_inv", + source=[["geo_int_jacobtot_inv"]], + fetch_func=[ff.fetch_s0c0], + label=r"$(J B)^{-1}$", + is_geo=True) +gk_quant_registry.register(_geo_int_jacobtot_inv) + +_geo_int_bmag = GkQuantity(name="geo_int_bmag", + source=[["geo_int_bmag"]], + fetch_func=[ff.fetch_s0c0], + label=r"$B$ (T)", + is_geo=True) +gk_quant_registry.register(_geo_int_bmag) + +# ----------------------------------------- vector geometric quantities (geo) +_geo_int_b_i = GkQuantity(name="geo_int_b_i", + source=[["geo_int_b_i"]], + fetch_func=[ff.fetch_s0cAll], + label=r"$b_%s$", + is_vector=True, + is_geo=True) +gk_quant_registry.register(_geo_int_b_i) + +# ------------------------------------------------------------------- field +_field = GkQuantity(name="field", + source=[["field"]], + fetch_func=[ff.fetch_s0c0], + label=r"$\phi$ (V)", + is_time_dep=True) +gk_quant_registry.register(_field) + +# --------------------------------------------------- plasma moments (per-sp) +_M0 = GkQuantity(name="M0", + source=[["M0"], ["M0M1M2"], ["M0M1M2parM2perp"], + ["MaxwellianMoments"], ["BiMaxwellianMoments"], + ["HamiltonianMoments"]], + fetch_func=[ff.fetch_s0c0] * 6, + label=r"$M_{0%s}$ (m$^{-3}$)", + is_species_dep=True, + is_time_dep=True) +gk_quant_registry.register(_M0) + +_M1 = GkQuantity(name="M1", + source=[["M1"], ["M0M1M2"], ["M0M1M2parM2perp"], + ["MaxwellianMoments"], ["BiMaxwellianMoments"], + ["HamiltonianMoments"]], + fetch_func=[ + ff.fetch_s0c0, ff.fetch_s0c1, ff.fetch_s0c1, + ff.fetch_s0c0_mul_s0c1, ff.fetch_s0c0_mul_s0c1, + ff.fetch_M1_from_H + ], + label=r"$M_{1%s}$ (m$^{-2}$/s)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M1) + +_M2par = GkQuantity( + name="M2par", + source=[["M2par"], ["M0M1M2parM2perp"], ["M2", "M2perp"]], + fetch_func=[ff.fetch_s0c0, ff.fetch_s0c2, ff.fetch_s0c0_sub_s1c0], + label=r"$M_{2\parallel%s}$ (m$^{-1}$/s$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M2par) + +_M2perp = GkQuantity( + name="M2perp", + source=[["M2perp"], ["M0M1M2parM2perp"], ["M2", "M2par"]], + fetch_func=[ff.fetch_s0c0, ff.fetch_s0c3, ff.fetch_s0c0_sub_s1c0], + label=r"$M_{2\perp%s}$ (m$^{-1}$/s$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M2perp) + +_M2 = GkQuantity(name="M2", + source=[["M2"], ["M0M1M2"], ["M0M1M2parM2perp"], + [_M2par, _M2perp]], + fetch_func=[ + ff.fetch_s0c0, ff.fetch_s0c2, ff.fetch_s0c2_add_s0c3, + ff.fetch_s0c0_add_s1c0 + ], + label=r"$M_{2%s}$ (m$^{-1}$/s$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M2) + +_M3par = GkQuantity(name="M3par", + source=[["M3par"]], + fetch_func=[ff.fetch_s0c0], + label=r"$M_{3\parallel%s}$ (1/s$^3$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M3par) + +_M3perp = GkQuantity(name="M3perp", + source=[["M3perp"]], + fetch_func=[ff.fetch_s0c0], + label=r"$M_{3\perp%s}$ (1/s$^3$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M3perp) + +_M3 = GkQuantity(name="M3", + source=[["M3"], [_M3par, _M3perp]], + fetch_func=[ff.fetch_s0c0, ff.fetch_s0c0_add_s1c0], + label=r"$M_{3%s}$ (1/s$^3$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_M3) + +_upar = GkQuantity( + name="upar", + source=[["MaxwellianMoments"], ["BiMaxwellianMoments"], [_M0, _M1]], + fetch_func=[ff.fetch_s0c1, ff.fetch_s0c1, ff.fetch_s1c0_div_s0c0], + label=r"$u_{\parallel %s}$ (m/s)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_upar) + +_Tpar = GkQuantity( + name="Tpar", + source=[["BiMaxwellianMoments"], [_M0, _M1, _M2par]], + fetch_func=[ff.fetch_Tpar_from_BiMax, ff.fetch_Tpar_from_M0_M1_M2par], + label=r"$T_{\parallel %s}$ (J)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_Tpar) + +_Tperp = GkQuantity( + name="Tperp", + source=[["BiMaxwellianMoments"], [_M0, _M2perp]], + fetch_func=[ff.fetch_Tperp_from_BiMax, ff.fetch_Tperp_from_M0_M2perp], + label=r"$T_{\perp %s}$ (J)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_Tperp) + +# ------------------------------------------- combined plasma moments (per-sp) +_temp = GkQuantity( + name="temp", + source=[["MaxwellianMoments"], [_Tpar, _Tperp]], + fetch_func=[ff.fetch_temp_from_Max, ff.fetch_temp_from_Tpar_Tperp], + label=r"$T_{%s}$ (J)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_temp) + +_press = GkQuantity(name="press", + source=[["MaxwellianMoments"], ["BiMaxwellianMoments"], + [_M0, _temp]], + fetch_func=[ + ff.fetch_press_from_Max, ff.fetch_press_from_BiMax, + ff.fetch_s0c0_mul_s1c0 + ], + label=r"$p_{%s}$ (Pa)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_press) + +_presspar = GkQuantity(name="presspar", + source=[[_M0, _Tpar]], + fetch_func=[ff.fetch_press_p], + label=r"$p_{\parallel %s}$ (Pa)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_presspar) + +_pressperp = GkQuantity(name="pressperp", + source=[[_M0, _Tperp]], + fetch_func=[ff.fetch_press_p], + label=r"$p_{\perp %s}$ (Pa)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_pressperp) + +_beta = GkQuantity(name="beta", + source=[[_geo_int_bmag, _press]], + fetch_func=[ff.fetch_beta_from_bmag_press], + label=r"$\beta_{%s}$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_beta) + +# --------------------------------------------------------------- heat fluxes +_qpar = GkQuantity(name="qpar", + source=[[_M3par]], + fetch_func=[ff.fetch_qpar], + label=r"$q_{\parallel %s}$ (W/m$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qpar) + +_qperp = GkQuantity(name="qperp", + source=[[_M3perp]], + fetch_func=[ff.fetch_qperp], + label=r"$q_{\perp %s}$ (W/m$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qperp) + +_qpar_fluid = GkQuantity(name="qpar_fluid", + source=[[_M0, _M1, _M2par, _M3par]], + fetch_func=[ff.fetch_qpar_fluid], + label=r"$q_{\parallel %s}^{fluid}$ (W/m$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qpar_fluid) + +_qperp_fluid = GkQuantity(name="qperp_fluid", + source=[[_M0, _M1, _M2perp, _M3perp]], + fetch_func=[ff.fetch_qperp_fluid], + label=r"$q_{\perp %s}^{fluid}$ (W/m$^2$)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qperp_fluid) + +# ------------------------------------------------ thermal speed / lengths +_vt = GkQuantity(name="vt", + source=[[_temp]], + fetch_func=[ff.fetch_vt], + label=r"$v_{t,%s}$ (m/s)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_vt) + +_larmor_radius = GkQuantity(name="larmor_radius", + source=[[_temp, _geo_int_bmag]], + fetch_func=[ff.fetch_larmor_radius], + label=r"$\rho_{%s}$ (m)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_larmor_radius) + +_debye_length = GkQuantity(name="debye_length", + source=[[_temp, _M0]], + fetch_func=[ff.fetch_debye_length], + label=r"$\lambda_{D,%s}$ (m)", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_debye_length) + +_c_s = GkQuantity(name="c_s", + source=[[_M0, _temp]], + fetch_func=[ff.fetch_c_s], + label=r"$c_{s}$ (m/s)", + is_time_dep=True, + is_species_dep=False, + is_multi_species=True) +gk_quant_registry.register(_c_s) + +# ----------------------------------------------------------- drift speeds +_ExB_vel = GkQuantity( + name="ExB_vel", + source=[[_geo_int_jacobtot_inv, _geo_int_bmag, _geo_int_b_i, _field]], + fetch_func=[ff.fetch_ExB_vel], + label=r"$v_{E,%s}$ (m/s)", + is_time_dep=True, + is_vector=True) +gk_quant_registry.register(_ExB_vel) + +_gradB_vel = GkQuantity( + name="gradB_vel", + source=[[_geo_int_jacobtot_inv, _geo_int_bmag, _geo_int_b_i, _Tperp]], + fetch_func=[ff.fetch_gradB_vel], + label=r"$v_{\nabla B,%s}$ (m/s)", + is_time_dep=True, + is_species_dep=True, + is_vector=True) +gk_quant_registry.register(_gradB_vel) + +_diamag_vel = GkQuantity(name="diamag_vel", + source=[[ + _geo_int_jacobtot_inv, _geo_int_bmag, _geo_int_b_i, + _M0, _pressperp + ]], + fetch_func=[ff.fetch_diamag_vel], + label=r"$v_{dia,%s}$ (m/s)", + is_time_dep=True, + is_species_dep=True, + is_vector=True) +gk_quant_registry.register(_diamag_vel) + +# ------------------------------------------------------------- phase space +_distf = GkQuantity(name="distf", + source=[[""]], + fetch_func=[ff.load_distf], + label=r"$f_{%s}$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_distf) + +# ----------------------------------------------------------- normalized +_rho_over_lambda = GkQuantity(name="rho_over_lambda", + source=[[_larmor_radius, _debye_length]], + fetch_func=[ff.fetch_rho_over_lambda], + label=r"$(\rho/\lambda_D)_{%s}$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_rho_over_lambda) + +_phi_norm = GkQuantity(name="phi_norm", + source=[[_field, _temp]], + fetch_func=[ff.fetch_phi_norm], + label=r"$e\phi/T_{%s}$", + is_time_dep=True, + is_species_dep=False) +gk_quant_registry.register(_phi_norm) + +_qpar_norm = GkQuantity(name="qpar_norm", + source=[[_qpar, _M0, _temp, _vt]], + fetch_func=[ff.fetch_qpar_norm], + label=r"$q_{\parallel %s}/(n T v_{th})$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qpar_norm) + +_qperp_norm = GkQuantity(name="qperp_norm", + source=[[_qperp, _M0, _temp, _vt]], + fetch_func=[ff.fetch_qperp_norm], + label=r"$q_{\perp %s}/(n T v_{th})$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qperp_norm) + +_qpar_fluid_norm = GkQuantity(name="qpar_fluid_norm", + source=[[_qpar_fluid, _M0, _temp, _vt]], + fetch_func=[ff.fetch_qpar_norm], + label=r"$q_{\parallel %s}^{fluid}/(n T v_{t})$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qpar_fluid_norm) + +_qperp_fluid_norm = GkQuantity(name="qperp_fluid_norm", + source=[[_qperp_fluid, _M0, _temp, _vt]], + fetch_func=[ff.fetch_qperp_norm], + label=r"$q_{\perp %s}^{fluid}/(n T v_{t})$", + is_time_dep=True, + is_species_dep=True) +gk_quant_registry.register(_qperp_fluid_norm) diff --git a/src/postgkyl/diagnostics/gk/rz.py b/src/postgkyl/diagnostics/gk/rz.py new file mode 100644 index 00000000..169d030f --- /dev/null +++ b/src/postgkyl/diagnostics/gk/rz.py @@ -0,0 +1,59 @@ +"""Compatibility aliases for the gyrokinetic R-Z operation. + +Canonical imports live in :mod:`postgkyl.operations.gyrokinetics`. This +module remains for the current major version and contains no copied +algorithm or defaults. +""" + +from postgkyl.operations.gyrokinetics.rz import ( + Geometry, + RzProjection, + geometry_prefix, + gk_rz, + map_to_rz, + per_block_path, + resolve_geometry, + resolve_rz_projection, +) + + +def rz_projections(datasets, + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + z_axis: float = 0.0, + nz_interp: int = 8) -> dict: + """Compatibility batch wrapper using this module's patchable aliases.""" + projections = {} + for data in datasets: + key = geometry_prefix(data.file_name) + if key in projections: + continue + block = data.ctx.get("block") + geometry = resolve_geometry(data.file_name, + mapc2p=per_block_path(mapc2p, block), + nodes_file=per_block_path(nodes_file, block)) + projections[key] = resolve_rz_projection(data, + geometry, + z_axis=z_axis, + nz_interp=nz_interp) + return projections + + +def projection_for(projections: dict, data): + """Return the compatibility projection belonging to ``data``'s block.""" + return projections[geometry_prefix(data.file_name)] + + +__all__ = [ + "Geometry", + "RzProjection", + "geometry_prefix", + "gk_rz", + "map_to_rz", + "per_block_path", + "projection_for", + "resolve_geometry", + "resolve_rz_projection", + "rz_projections", +] diff --git a/src/postgkyl/diagnostics/gk/utils.py b/src/postgkyl/diagnostics/gk/utils.py new file mode 100644 index 00000000..5abffd25 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/utils.py @@ -0,0 +1,223 @@ +"""Small file helpers shared by the gyrokinetic loaders and the +layer-13 program-scale diagnostics. + +Ported from ``src_bak/postgkyl/gk/gk_utils.py``. ``read_gfile``/ +``read_interpolated_gfile`` are adapted to the new API (``postgkyl.gdata.load`` ++ ``.interpolate()``) in place of the retired ``GData``/``GInterpModal`` pair; +``read_gfile_if_present`` drops the old code's ``verb_print(ctx, ...)`` call +(``ctx`` was never a parameter of that function in ``src_bak`` -- an existing +bug -- and printing belongs to the CLI, not a loader) in favor of returning a +plain ``found`` flag. ``read_time_trace_if_present`` and +``set_tick_font_size`` are shared by the three program-scale diagnostics that +build figures directly with matplotlib (``energy_balance``, +``particle_balance``, ``nodes``) rather than each keeping its own private +copy. +""" + +from __future__ import annotations + +import glob +import os + +import numpy as np +from postgkyl import numerics +from postgkyl.gdata import GData + +# Maximum number of blocks a multiblock simulation is assumed to have, used +# only to bound an open-ended slice request in get_block_indices. +MAX_NUM_BLOCKS = 10000 + + +def read_gfile(file_name: str) -> tuple[list[np.ndarray], np.ndarray, GData]: + """Read a Gkeyll file, squeezing singleton axes out of the grid and values. + + Args: + file_name: Path to the ``.gkyl`` file. + + Returns: + ``(grid, values, gdata)``: the squeezed grid (a list of squeezed + per-dimension arrays -- ``GDataState.grid`` never hands back a bare + ``ndarray``), the squeezed value array, and the loaded dataset itself + (for further chaining). + """ + gdata = GData(file_name) + grid = gdata.get_grid() + values = gdata.get_values() + grid_out = [np.squeeze(grid[d]) for d in range(len(grid))] + return grid_out, np.squeeze(values), gdata + + +def read_gfile_if_present( + file_name: str, +) -> tuple[bool, list[np.ndarray] | None, np.ndarray | None, GData | None]: + """Read a Gkeyll file if it exists. + + Args: + file_name: Path to the file. + + Returns: + ``(found, grid, values, gdata)``; ``found`` is False and the remaining + entries are ``None`` when ``file_name`` does not exist. + """ + if not os.path.exists(file_name): + return False, None, None, None + grid, values, gdata = read_gfile(file_name) + return True, grid, values, gdata + + +def read_time_trace_if_present( + file_name: str, +) -> tuple[bool, np.ndarray | None, np.ndarray | None, GData | None]: + """Read a 1-D time-trace file if present: ``(found, time, values, gdata)``. + + ``read_gfile_if_present`` always returns the grid as a *list* of + per-dimension arrays (``GDataState.grid`` never hands back a bare + ``ndarray``, only a list of one for 1-D data) -- this unwraps that single + entry into the plain time array every trace in + :mod:`~postgkyl.diagnostics.gk.energy_balance`/ + :mod:`~postgkyl.diagnostics.gk.particle_balance` is indexed + against. + """ + found, grid, values, gdata = read_gfile_if_present(file_name) + time = grid[0] if found else None + return found, time, values, gdata + + +def read_interpolated_gfile( + file_name: str, + poly_order: int, + basis_type: str, + comp: int | str | None = None, +) -> tuple[list[np.ndarray], np.ndarray, GData]: + """Read a Gkeyll file and interpolate it onto a uniform mesh. + + Args: + file_name: Path to the file. + poly_order: Polynomial order of the DG basis. + basis_type: Long basis name, e.g. ``"serendipity"``. + comp: Optional component selector applied *after* interpolation + (an int index or a ``"start:stop"`` slice string); ``None`` keeps + every component. + + Returns: + ``(grid, values, gdata)``: the squeezed interpolated grid (a list of + squeezed per-dimension arrays) and values, and the interpolated dataset. + """ + gdata = GData(file_name, basis_type=basis_type, poly_order=poly_order) + interpolated = gdata.interpolate() + if comp is not None: + interpolated = interpolated.select(comp=comp) + grid = interpolated.get_grid() + values = interpolated.get_values() + grid_out = [np.squeeze(grid[d]) for d in range(len(grid))] + return grid_out, np.squeeze(values), interpolated + + +def interpolated_grid_values( + data: GData, + *, + comp: int = 0) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]: + """Interpolate ``data``'s DG coefficients onto its computational mesh. + + Shared by :mod:`~postgkyl.diagnostics.gk.rz` and + :mod:`~postgkyl.diagnostics.gk.fluxsurf`, which both need a + field-aligned dataset's fine computational grid (for sampling the + simulation's geometry) alongside its interpolated values. + + Args: + data: The dataset. Normally not yet interpolated; an already-interpolated + dataset is accepted and used as-is. + comp: Component to return. + + Returns: + ``(edges, centers, values)``: the refined edge grid (one 1-D array per + dimension), its cell-centered equivalent, and component ``comp``'s + values on that grid. + """ + # Idempotent on purpose: 'pgkyl ... interp gk_rz' is a natural thing to + # type, and interpolating twice would run the DG evaluation matrix over + # values that are already point values -- silently wrong output rather + # than an error. + field = data if data.ctx.get("interpolated") else data.interpolate() + cells = field.values.shape[:-1] + centers = numerics.nodal_to_cell_centered_grid(field.grid, cells) + return field.grid, centers, field.values[..., comp] + + +def set_tick_font_size(ax, size: float) -> None: + """Set an axes' tick-label and offset-text font size to ``size``.""" + ax.tick_params(axis="both", labelsize=size) + ax.yaxis.get_offset_text().set_size(size) + ax.xaxis.get_offset_text().set_size(size) + + +def dict_get_bool(dict_in: dict, key: str, default: bool) -> bool: + """Interpret a dict value as a bool, returning ``default`` if absent. + + String values ``'1'``/``'true'`` (case-insensitive) are True, anything + else False; non-string values are converted with ``bool()``. + """ + if key not in dict_in: + return default + val = dict_in[key] + if isinstance(val, str): + return val.strip().lower() in ("1", "true") + return bool(val) + + +def parse_slice_string(value: str) -> slice: + """Parse a ``slice()`` from a ``'start:stop:step'`` string. + + Raises: + ValueError: if any non-empty part is not an integer. + """ + parts = value.split(":") + parsed_parts = [] + for p in parts: + try: + parsed_parts.append(int(p) if p else None) + except ValueError: + raise ValueError(f"Invalid slice part: {p}") + return slice(*parsed_parts) + + +def get_block_indices(multib: str, file_path_name: str) -> list[int]: + """Return the indices of the blocks to process in a multiblock simulation. + + Args: + multib: ``"-10"`` for a single block (index 0); ``"-1"`` to discover and + use every block found by globbing ``file_path_name``; otherwise a + comma-separated list or a ``'start:stop[:step]'`` slice string of the + desired block indices. + file_path_name: Path/filename glob used to discover blocks when + ``multib == "-1"``, with the block index replaced by ``"*"`` (e.g. + ``"_b*-_field_0.gkyl"``). + + Returns: + A list of block indices. + + Raises: + NameError: if ``multib`` is neither ``"-10"``/``"-1"``, a comma-separated + list, a slice string, nor a single integer. + """ + + def _is_int(s: str) -> bool: + try: + int(s) + return True + except ValueError: + return False + + if multib == "-10": + return [0] + if multib == "-1": + return list(range(len(glob.glob(file_path_name)))) + if "," in multib: + return [int(b) for b in multib.split(",")] + if ":" in multib: + s = parse_slice_string(multib) + return list(range(*s.indices(MAX_NUM_BLOCKS))) + if _is_int(multib): + return [int(multib)] + raise NameError( + "Blocks given to --multib -m must be a comma separated list or slice.") diff --git a/src/postgkyl/diagnostics/mom/__init__.py b/src/postgkyl/diagnostics/mom/__init__.py new file mode 100644 index 00000000..840dac64 --- /dev/null +++ b/src/postgkyl/diagnostics/mom/__init__.py @@ -0,0 +1,23 @@ +"""Diagnostics for Gkeyll fluid-moment equation systems.""" + +from . import ( + enstrophy, + five_moment, + ke_dke, + mhd, + multispecies, + plasma, + rotations, + ten_moment, +) + +__all__ = [ + "enstrophy", + "five_moment", + "ke_dke", + "mhd", + "multispecies", + "plasma", + "rotations", + "ten_moment", +] diff --git a/src/postgkyl/diagnostics/mom/enstrophy.py b/src/postgkyl/diagnostics/mom/enstrophy.py new file mode 100644 index 00000000..aa78ce22 --- /dev/null +++ b/src/postgkyl/diagnostics/mom/enstrophy.py @@ -0,0 +1,126 @@ +"""2-D/3-D five-moment enstrophy diagnostic. + +Ported from ``src_bak/postgkyl/tools/calc_enstrophy.py``. Sweeps a family of +five-moment output frames (density + momentum, ``rho, px, py, pz``) and +computes, per frame, the enstrophy in its general form (integral of the +squared magnitude of the curl of the velocity over the volume) and its +incompressible form (integral of a velocity-gradient invariant, weighted by +density). + +Fixes one bug present in ``src_bak``: ``incom_enstrophy = enstrophy`` aliased +the very array the general-form result was written into, so both returned +traces ended up identical (equal to whichever form was written last in the +frame loop) instead of being the two distinct quantities the function's own +docstring and return statement promised -- doctrine #21 requires fixing an +unambiguous bug rather than silently porting it forward. The per-cell nested +loop's ``range(len(axis) - 1)`` bound (leaving the last plane along every +axis at zero) is preserved verbatim: unlike the aliasing, it is not +unambiguously a bug (it could be deliberate avoidance of a less-accurate +``np.gradient`` edge-order boundary), so changing it would be a silent +numerical-behavior change doctrine #21 forbids. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from postgkyl.gdata import GData + + +@dataclass(frozen=True) +class EnstrophyTraces: + """Per-frame enstrophy traces, one entry per swept frame. + + Attributes: + enstrophy: General-form enstrophy (integral of the squared curl + magnitude). + incompressible_enstrophy: Incompressible-form enstrophy (integral of a + density-weighted velocity-gradient invariant). + """ + + enstrophy: np.ndarray + incompressible_enstrophy: np.ndarray + + +def _enstrophy_terms(rho: np.ndarray, px: np.ndarray, py: np.ndarray, + pz: np.ndarray, dx: float, dy: float, + dz: float) -> tuple[float, float]: + """Pure array math: the general and incompressible enstrophy integrals + for one frame of five-moment (density + momentum) data. + + Args: + rho, px, py, pz: 3-D density and momentum-component arrays (same shape). + dx, dy, dz: Grid spacing along each axis. + + Returns: + ``(enstrophy, incompressible_enstrophy)``: the two scalar integrals for + this frame. + """ + u = px / rho + v = py / rho + w = pz / rho + + u_grad = np.gradient(u, dx, dy, dz, edge_order=2) + v_grad = np.gradient(v, dx, dy, dz, edge_order=2) + w_grad = np.gradient(w, dx, dy, dz, edge_order=2) + grad_tensor = np.array([u_grad, v_grad, w_grad]) + + u_x, u_y, u_z = u_grad + v_x, v_y, v_z = v_grad + w_x, w_y, w_z = w_grad + + curl_mag = (w_y - v_z)**2 + (u_z - w_x)**2 + (v_x - u_y)**2 + enstrophy = np.sum(curl_mag, axis=(0, 1, 2)) * dx * dy * dz + + nx, ny, nz = rho.shape + incom_mag = np.zeros((nx, ny, nz)) + for c in range(nx - 1): + for j in range(ny - 1): + for k in range(nz - 1): + cell = grad_tensor[:, :, c, j, k] + incom_mag[c, j, k] = np.trace(np.transpose(cell) * cell) * rho[c, j, k] + incompressible_enstrophy = np.sum(incom_mag, axis=(0, 1, 2)) * dx * dy * dz + + return enstrophy, incompressible_enstrophy + + +def enstrophy( + stem: str, + init_frame: int, + final_frame: int, + *, + extension: str = "gkyl", +) -> EnstrophyTraces: + """Sweep a frame family and compute the enstrophy in 2 forms. + + Args: + stem: File-name stem before the frame number, e.g. ``"sim-fluid_"``. + init_frame: First frame (inclusive). + final_frame: Last frame (inclusive). + extension: File extension of the frame files (defaults to the native + ``gkyl`` format). + + Returns: + :class:`EnstrophyTraces`, one entry per swept frame. + """ + num_frames = final_frame - init_frame + 1 + + first = GData(f"{stem}{init_frame}.{extension}") + grid = first.grid + dx = grid[0][1] - grid[0][0] + dy = grid[1][1] - grid[1][0] + dz = grid[2][1] - grid[2][0] + + enstrophy_trace = np.empty(num_frames) + incompressible_trace = np.empty(num_frames) + for r, frame_idx in enumerate(range(init_frame, final_frame + 1)): + data = GData(f"{stem}{frame_idx}.{extension}") + values = data.values + rho, px, py, pz = (values[..., c] for c in range(4)) + enstrophy_trace[r], incompressible_trace[r] = _enstrophy_terms( + rho, px, py, pz, dx, dy, dz) + + return EnstrophyTraces(enstrophy=enstrophy_trace, + incompressible_enstrophy=incompressible_trace) diff --git a/src/postgkyl/diagnostics/mom/five_moment.py b/src/postgkyl/diagnostics/mom/five_moment.py new file mode 100644 index 00000000..d726a5c6 --- /dev/null +++ b/src/postgkyl/diagnostics/mom/five_moment.py @@ -0,0 +1,556 @@ +"""Five-moment (Euler) diagnostics -- density, velocity, pressure, temperature, +sound speed, Mach number. + +Fluid moment data is laid out ``[rho, rho*vx, rho*vy, rho*vz, E, ...]``: the +first four components are shared with 10-moment/MHD data, and ``pressure``/ +``ke``/``temp``/``sound``/``mach`` additionally accept 10-moment data +(``num_moms=10``), inferring which layout applies from the number of +components when ``num_moms`` is not given. + +Each public function takes a ``GDataState`` and returns one (funneling +through ``_result``); the array-level math is kept in module-private +``_get_*`` helpers, copied verbatim from the pre-restructure ``models`` / +``operations`` layers (06/08) so ``ten_moment``/``mhd``/``plasma``/``multispecies`` +can compose the same formulas without re-deriving them. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ...gdatastate.guards import require_field_domain as _require_field_domain + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = ("extracting primitive variables from raw DG coefficients would " + "mix basis functions") + + +# --------------------------------------------------------- array-level math +def _get_density(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the (mass) density from fluid moment data. + + The density is component 0 of the moment array. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension. + values: Moment array whose last axis holds the conserved variables. + + Returns: + ``(grid, values)`` with the density as a single trailing component. + """ + return list(grid), values[..., 0, np.newaxis] + + +def _get_vx(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the x velocity: x momentum (component 1) over density.""" + _, rho = _get_density(grid, values) + return list(grid), values[..., 1, np.newaxis] / rho + + +def _get_vy(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the y velocity: y momentum (component 2) over density.""" + _, rho = _get_density(grid, values) + return list(grid), values[..., 2, np.newaxis] / rho + + +def _get_vz(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the z velocity: z momentum (component 3) over density.""" + _, rho = _get_density(grid, values) + return list(grid), values[..., 3, np.newaxis] / rho + + +def _get_vi(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the velocity vector ``(vx, vy, vz)``: momentum (1:4) over density.""" + _, rho = _get_density(grid, values) + return list(grid), values[..., 1:4] / rho + + +def _infer_num_moms(values: np.ndarray, num_moms: int | None) -> int: + """Resolve the moment count, inferring it from the component count.""" + if num_moms is not None: + return num_moms + num_comps = values.shape[-1] + if num_comps == 5: + return 5 + if num_comps == 10: + return 10 + raise ValueError( + f"Number of components appears to be {num_comps:d}; it needs to be " + "specified using 'num_moms' (5 or 10)") + + +def _get_p( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the scalar pressure from fluid moment data. + + For 5-moment data the pressure is the total energy minus the bulk kinetic + energy, scaled by ``gas_gamma - 1``. For 10-moment data it is the trace of + the pressure tensor over three: ``(P_xx + P_yy + P_zz) / 3``. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension. + values: Moment array (5- or 10-moment). + gas_gamma: Adiabatic index, used only for 5-moment data. + num_moms: Number of moments (5 or 10); inferred from the component count + when ``None``. + + Returns: + ``(grid, values)`` holding the scalar pressure field. + + Raises: + ValueError: If ``num_moms`` is ``None`` and cannot be inferred. + """ + num_moms = _infer_num_moms(values, num_moms) + + if num_moms == 5: + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + out_values = (gas_gamma - 1) * (values[..., 4, np.newaxis] - 0.5 * rho * + (vx**2 + vy**2 + vz**2)) + else: # num_moms == 10 + # Trace of the pressure tensor, computed inline (rather than calling + # ten_moment._get_pxx/_get_pyy/_get_pzz) to keep five_moment -> + # ten_moment a one-way edge; ten_moment._get_pxx/pyy/pzz apply this same + # M_ii - rho*v_i*v_i formula component-wise. + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + pxx = values[..., 4, np.newaxis] - rho * vx * vx + pyy = values[..., 7, np.newaxis] - rho * vy * vy + pzz = values[..., 9, np.newaxis] - rho * vz * vz + out_values = (pxx + pyy + pzz) / 3.0 + + return list(grid), out_values + + +def _get_ke( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the kinetic (bulk-flow) energy density from fluid moment data. + + For 5-moment data it is the total energy minus the thermal energy + ``p / (gas_gamma - 1)``. For 10-moment data it is + ``0.5 * rho * (vx**2 + vy**2 + vz**2)`` directly. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension. + values: Moment array (5- or 10-moment). + gas_gamma: Adiabatic index, used only for 5-moment data. + num_moms: Number of moments (5 or 10); inferred from the component count + when ``None``. + + Returns: + ``(grid, values)`` holding the kinetic energy density field. + """ + num_moms = _infer_num_moms(values, num_moms) + + if num_moms == 5: + _, pr = _get_p(grid, values, gas_gamma=gas_gamma, num_moms=num_moms) + out_values = values[..., 4, np.newaxis] - pr / (gas_gamma - 1) + else: # num_moms == 10 + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + out_values = 0.5 * rho * (vx**2 + vy**2 + vz**2) + + return list(grid), out_values + + +def _get_temp( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the temperature ``T = p / rho`` from fluid moment data.""" + _, rho = _get_density(grid, values) + _, pr = _get_p(grid, values, gas_gamma=gas_gamma, num_moms=num_moms) + return list(grid), pr / rho + + +def _get_sound( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the sound speed ``c_s = sqrt(gas_gamma * p / rho)``.""" + _, rho = _get_density(grid, values) + _, pr = _get_p(grid, values, gas_gamma=gas_gamma, num_moms=num_moms) + return list(grid), np.sqrt(gas_gamma * pr / rho) + + +def _get_mach( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the sonic Mach number ``M = |v| / c_s``.""" + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + _, cs = _get_sound(grid, values, gas_gamma=gas_gamma, num_moms=num_moms) + return list(grid), np.sqrt(vx**2 + vy**2 + vz**2) / cs + + +# ---------------------------------------------------------------- GData verbs +def density(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Mass density (component 0 of fluid moment data). + + Args: + data: Fluid moment data; must be NumPy-backed. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the density. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "density", _REASON) + grid, values = _get_density(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def xvel(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """x velocity: x momentum (component 1) over density. + + Args: + data: Fluid moment data; must be NumPy-backed. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the x velocity. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "xvel", _REASON) + grid, values = _get_vx(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def yvel(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """y velocity: y momentum (component 2) over density. + + Args: + data: Fluid moment data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the y velocity. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "yvel", _REASON) + grid, values = _get_vy(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def zvel(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """z velocity: z momentum (component 3) over density. + + Args: + data: Fluid moment data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the z velocity. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "zvel", _REASON) + grid, values = _get_vz(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def vel(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Velocity vector ``(vx, vy, vz)``: momentum (1:4) over density. + + Args: + data: Fluid moment data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A three-component dataset of the fluid velocity. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "vel", _REASON) + grid, values = _get_vi(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pressure(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Scalar pressure from fluid moment data (5- or 10-moment). + + Args: + data: Fluid moment data (5- or 10-moment); must be NumPy-backed. + gas_gamma: Adiabatic index, used only for 5-moment data. + num_moms: Number of moments (5 or 10); inferred from the component count + when ``None``. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the scalar pressure. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or ``num_moms`` is + ``None`` and cannot be inferred. + """ + _require_field_domain(data, "pressure", _REASON) + grid, values = _get_p(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=num_moms) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def ke(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Kinetic (bulk-flow) energy density from fluid moment data. + + Args: + data: Fluid moment data; must be NumPy-backed. + gas_gamma: Adiabatic index, used only for 5-moment data. + num_moms: Number of moments (5 or 10); inferred from component count + when ``None``. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the bulk-flow energy density. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or ``num_moms`` is + ``None`` and cannot be inferred. + """ + _require_field_domain(data, "ke", _REASON) + grid, values = _get_ke(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=num_moms) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def temp(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Temperature ``T = p / rho`` from fluid moment data. + + Args: + data: Fluid moment data; must be NumPy-backed. + gas_gamma: Adiabatic index, used only for 5-moment data. + num_moms: Number of moments (5 or 10); inferred from component count + when ``None``. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the temperature. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or ``num_moms`` is + ``None`` and cannot be inferred. + """ + _require_field_domain(data, "temp", _REASON) + grid, values = _get_temp(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=num_moms) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def sound(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Sound speed ``c_s = sqrt(gas_gamma * p / rho)``. + + Args: + data: Fluid moment data; must be NumPy-backed. + gas_gamma: Adiabatic index. + num_moms: Number of moments (5 or 10); inferred from component count + when ``None``. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the sound speed. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or ``num_moms`` is + ``None`` and cannot be inferred. + """ + _require_field_domain(data, "sound", _REASON) + grid, values = _get_sound(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=num_moms) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def mach(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Sonic Mach number ``M = |v| / c_s``. + + Args: + data: Fluid moment data; must be NumPy-backed. + gas_gamma: Adiabatic index used to compute the sound speed. + num_moms: Number of moments (5 or 10); inferred from component count + when ``None``. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the Mach number. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or ``num_moms`` is + ``None`` and cannot be inferred. + """ + _require_field_domain(data, "mach", _REASON) + grid, values = _get_mach(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=num_moms) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def velocity(density: "GDataState", + momentum: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Velocity from separate density and momentum moments. + + Computes the flow velocity by dividing the ``momentum`` moments by the + ``density`` moment, component-wise. The two inputs are assumed to share + the same grid; the result carries the ``density`` dataset's grid. + + Args: + density: Number/mass density moment (single component); the divisor. + Must be NumPy-backed. + momentum: Momentum moment(s) to divide by the density. Must be + NumPy-backed. + inplace: mutate and return ``density`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset of the velocity. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(density, "velocity", _REASON) + _require_field_domain(momentum, "velocity", _REASON) + values = momentum.values / density.values + return density._result(density.grid, + values, + inplace=inplace, + tag=tag, + label=label) + + +VARIABLES = { + "density": density, + "xvel": xvel, + "yvel": yvel, + "zvel": zvel, + "vel": vel, + "pressure": pressure, + "ke": ke, + "temp": temp, + "sound": sound, + "mach": mach, +} diff --git a/src/postgkyl/diagnostics/mom/ke_dke.py b/src/postgkyl/diagnostics/mom/ke_dke.py new file mode 100644 index 00000000..b73bc5b4 --- /dev/null +++ b/src/postgkyl/diagnostics/mom/ke_dke.py @@ -0,0 +1,117 @@ +"""Five-moment kinetic-energy / dissipation-rate diagnostic. + +Ported from ``src_bak/postgkyl/tools/calc_ke_dke.py``. Sweeps a family of +five-moment output frames (density + momentum, ``rho, px, py, pz``), +integrates the kinetic energy over the grid for each frame, and estimates +its dissipation rate by backward finite difference between consecutive +frames. + +Fixes three bugs present in ``src_bak`` (doctrine #21: fix an unambiguous +bug rather than silently port it forward): + + - the per-frame file name inside the sweep loop was built as + ``f"root_file_name{c:d}.gkyl"`` -- a literal string containing the + parameter's *name*, not an f-string interpolating its *value* + (``f"{root_file_name}{c:d}.gkyl"``); only the *first* frame, read once + before the loop to get the grid spacing, used the correct spelling; + - ``dEk = ke`` aliased the very array the kinetic-energy trace was + written into (instead of allocating its own array), so writing the + dissipation-rate trace corrupted not-yet-read kinetic-energy values; + - the difference loop's ``range(init_frame, final_frame - 1)`` is off by + one frame short of every valid backward difference (it should run + through ``final_frame - 1`` inclusive, i.e. ``range(init_frame, + final_frame)``). + +Combined, no variant of the original code could ever have produced a +meaningful trace, so this ports the clearly-intended calculation (every +consecutive-frame backward difference) rather than reproducing undefined +behavior. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from postgkyl.gdata import GData + + +@dataclass(frozen=True) +class KineticEnergyTraces: + """Per-frame kinetic-energy traces. + + Attributes: + ke: Integrated kinetic energy, one entry per swept frame. + dke: Dissipation rate (backward difference of ``ke``), one entry per + consecutive frame pair -- one shorter than ``ke``. + """ + + ke: np.ndarray + dke: np.ndarray + + +def _kinetic_energy(rho: np.ndarray, px: np.ndarray, py: np.ndarray, + pz: np.ndarray, dx: float, dy: float, dz: float, + vol: float) -> float: + """Pure array math: the integrated kinetic energy for one frame.""" + u = px / rho + v = py / rho + w = pz / rho + e = rho * (u**2 + v**2 + w**2) + return np.sum(e, axis=(0, 1, 2)) * dx * dy * dz * vol + + +def _dissipation_rate(ke: np.ndarray, dt: float) -> np.ndarray: + """Backward-difference dissipation rate between every consecutive pair: + ``dke[i] = -(ke[i + 1] - ke[i]) / dt``.""" + return -(ke[1:] - ke[:-1]) / dt + + +def ke_dke( + root_file_name: str, + init_frame: int, + final_frame: int, + dim: int, + vol: float, + init_time: float, + final_time: float, + *, + extension: str = "gkyl", +) -> KineticEnergyTraces: + """Sweep a frame family and compute the kinetic energy and dissipation rate. + + Args: + root_file_name: File-name stem before the frame number. + init_frame: First frame (inclusive). + final_frame: Last frame (inclusive). + dim: Simulation dimensionality (2 or 3); the z grid spacing is taken as + 1 when ``dim != 3``. + vol: Grid cell volume factor. + init_time: Simulation start time. + final_time: Simulation end time; used with ``init_time`` to derive a + uniform ``dt`` for the dissipation-rate estimate. + extension: File extension of the frame files (defaults to the native + ``gkyl`` format). + + Returns: + :class:`KineticEnergyTraces`. + """ + num_frames = final_frame - init_frame + 1 + dt = (final_time - init_time + 1) / num_frames + + first = GData(f"{root_file_name}{init_frame}.{extension}") + grid = first.grid + dx = grid[0][1] - grid[0][0] + dy = grid[1][1] - grid[1][0] + dz = grid[2][1] - grid[2][0] if dim == 3 else 1 + + ke = np.empty(num_frames) + for r, frame_idx in enumerate(range(init_frame, final_frame + 1)): + data = GData(f"{root_file_name}{frame_idx}.{extension}") + values = data.values + rho, px, py, pz = (values[..., c] for c in range(4)) + ke[r] = _kinetic_energy(rho, px, py, pz, dx, dy, dz, vol) + + dke = _dissipation_rate(ke, dt) + return KineticEnergyTraces(ke=ke, dke=dke) diff --git a/src/postgkyl/diagnostics/mom/mhd.py b/src/postgkyl/diagnostics/mom/mhd.py new file mode 100644 index 00000000..ca19013e --- /dev/null +++ b/src/postgkyl/diagnostics/mom/mhd.py @@ -0,0 +1,390 @@ +"""Ideal-MHD diagnostics -- the five-moment set (density/velocity) plus the +magnetic field, magnetic pressure, thermal pressure, temperature, sound +speed, and Mach number. + +MHD moment data is laid out ``[rho, mx, my, mz, E, Bx, By, Bz]``: components +0:4 are shared with the 5-moment layout (density and momentum), so density +and velocity are reused from :mod:`postgkyl.diagnostics.mom.five_moment`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ...gdatastate.guards import require_field_domain as _require_field_domain +from .five_moment import _get_density, _get_vx, _get_vy, _get_vz +from .five_moment import density, xvel, yvel, zvel, vel + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = ("extracting primitive variables from raw DG coefficients would " + "mix basis functions") + + +# --------------------------------------------------------- array-level math +def _get_mhd_Bx(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the x magnetic-field component (component 5 of MHD data).""" + return list(grid), values[..., 5, np.newaxis] + + +def _get_mhd_By(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the y magnetic-field component (component 6 of MHD data).""" + return list(grid), values[..., 6, np.newaxis] + + +def _get_mhd_Bz(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the z magnetic-field component (component 7 of MHD data).""" + return list(grid), values[..., 7, np.newaxis] + + +def _get_mhd_Bi(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Extract the magnetic-field vector ``(Bx, By, Bz)`` (components 5:8).""" + return list(grid), values[..., 5:8] + + +def _get_mhd_mag_p(grid: list[np.ndarray], + values: np.ndarray, + *, + mu_0: float = 1.0) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the magnetic pressure + ``p_B = 0.5 * (Bx**2 + By**2 + Bz**2) / mu_0``.""" + _, Bx = _get_mhd_Bx(grid, values) + _, By = _get_mhd_By(grid, values) + _, Bz = _get_mhd_Bz(grid, values) + return list(grid), 0.5 * (Bx**2 + By**2 + Bz**2) / mu_0 + + +def _get_mhd_p( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the thermal (gas) pressure. + + ``p = (gas_gamma - 1) * (E - 0.5*rho*|v|**2 - p_B)``. + """ + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + _, mag_p = _get_mhd_mag_p(grid, values, mu_0=mu_0) + + out_values = (gas_gamma - 1) * (values[..., 4, np.newaxis] - 0.5 * rho * + (vx**2 + vy**2 + vz**2) - mag_p) + return list(grid), out_values + + +def _get_mhd_temp( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the temperature ``T = p / rho``.""" + _, rho = _get_density(grid, values) + _, pr = _get_mhd_p(grid, values, gas_gamma=gas_gamma, mu_0=mu_0) + return list(grid), pr / rho + + +def _get_mhd_sound( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the sound speed ``c_s = sqrt(gas_gamma * p / rho)``.""" + _, rho = _get_density(grid, values) + _, pr = _get_mhd_p(grid, values, gas_gamma=gas_gamma, mu_0=mu_0) + return list(grid), np.sqrt(gas_gamma * pr / rho) + + +def _get_mhd_mach( + grid: list[np.ndarray], + values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the sonic Mach number ``M = |v| / c_s``.""" + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + _, cs = _get_mhd_sound(grid, values, gas_gamma=gas_gamma, mu_0=mu_0) + return list(grid), np.sqrt(vx**2 + vy**2 + vz**2) / cs + + +# ---------------------------------------------------------------- GData verbs +def bx(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """x magnetic-field component (component 5 of MHD data). + + Args: + data: MHD conserved variables; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``Bx``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "bx", _REASON) + grid, values = _get_mhd_Bx(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def by(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """y magnetic-field component (component 6 of MHD data). + + Args: + data: MHD conserved variables; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``By``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "by", _REASON) + grid, values = _get_mhd_By(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def bz(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """z magnetic-field component (component 7 of MHD data). + + Args: + data: MHD conserved variables; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``Bz``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "bz", _REASON) + grid, values = _get_mhd_Bz(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def bi(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Magnetic-field vector ``(Bx, By, Bz)`` (components 5:8). + + Args: + data: MHD conserved variables; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A three-component dataset of the magnetic field. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "bi", _REASON) + grid, values = _get_mhd_Bi(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def mag_pressure(data: "GDataState", + *, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Magnetic pressure ``p_B = 0.5 * (Bx**2 + By**2 + Bz**2) / mu_0``. + + Args: + data: MHD conserved variables; must be NumPy-backed. + mu_0: Vacuum permeability. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the magnetic pressure. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "mag_pressure", _REASON) + grid, values = _get_mhd_mag_p(data.grid, data.values, mu_0=mu_0) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pressure(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Thermal (gas) pressure + ``p = (gas_gamma - 1) * (E - 0.5*rho*|v|**2 - p_B)``. + + Args: + data: MHD conserved variables; must be NumPy-backed. + gas_gamma: Adiabatic index. + mu_0: Vacuum permeability used in the magnetic energy. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the thermal pressure. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pressure", _REASON) + grid, values = _get_mhd_p(data.grid, + data.values, + gas_gamma=gas_gamma, + mu_0=mu_0) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def temp(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Temperature ``T = p / rho``. + + Args: + data: MHD conserved variables; must be NumPy-backed. + gas_gamma: Adiabatic index. + mu_0: Vacuum permeability used to compute the pressure. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the temperature. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "temp", _REASON) + grid, values = _get_mhd_temp(data.grid, + data.values, + gas_gamma=gas_gamma, + mu_0=mu_0) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def sound(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Sound speed ``c_s = sqrt(gas_gamma * p / rho)``. + + Args: + data: MHD conserved variables; must be NumPy-backed. + gas_gamma: Adiabatic index. + mu_0: Vacuum permeability used to compute the pressure. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the sound speed. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "sound", _REASON) + grid, values = _get_mhd_sound(data.grid, + data.values, + gas_gamma=gas_gamma, + mu_0=mu_0) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def mach(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Sonic Mach number ``M = |v| / c_s``. + + Args: + data: MHD conserved variables; must be NumPy-backed. + gas_gamma: Adiabatic index. + mu_0: Vacuum permeability used to compute the pressure. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the Mach number. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "mach", _REASON) + grid, values = _get_mhd_mach(data.grid, + data.values, + gas_gamma=gas_gamma, + mu_0=mu_0) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +VARIABLES = { + "density": density, + "xvel": xvel, + "yvel": yvel, + "zvel": zvel, + "vel": vel, + "Bx": bx, + "By": by, + "Bz": bz, + "Bi": bi, + "magpressure": mag_pressure, + "pressure": pressure, + "temp": temp, + "sound": sound, + "mach": mach, +} diff --git a/src/postgkyl/diagnostics/mom/multispecies.py b/src/postgkyl/diagnostics/mom/multispecies.py new file mode 100644 index 00000000..f47f2e05 --- /dev/null +++ b/src/postgkyl/diagnostics/mom/multispecies.py @@ -0,0 +1,217 @@ +"""Multi-species diagnostics: energy-balance decomposition and current +accumulation. + +``energetics`` separates a two-species (electron + ion) fluid/field system +into its constituent energy components; ``accumulate_current`` scales a +single species' moment data by its charge (or charge-to-mass ratio) so that +several species can be summed into a total current. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ... import numerics +from ...gdatastate.guards import require_field_domain as _require_field_domain +from .five_moment import _get_ke, _get_p + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = "decomposing energy from raw DG coefficients would mix basis functions" + + +# --------------------------------------------------------- array-level math +def _energetics( + elc_grid: list[np.ndarray], + elc_values: np.ndarray, + ion_grid: list[np.ndarray], + ion_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Separate a two-species plasma's energy into its constituent parts. + + Args: + elc_grid: Electron moment grid. + elc_values: Electron fluid moment array. + ion_grid: Ion moment grid. + ion_values: Ion fluid moment array. + field_grid: EM field grid. + field_values: EM field array laid out ``[Ex, Ey, Ez, Bx, By, Bz]``. + gas_gamma: Adiabatic index, forwarded to the pressure/kinetic-energy + calculation for both species. + num_moms: Number of moments (5 or 10) for both species; inferred from + the component count when ``None``. + + Returns: + ``(grid, values)`` with a 7-component field: + ``(electron thermal, electron kinetic, ion thermal, ion kinetic, + electric, magnetic, total)``. + """ + out = np.zeros(field_values.shape[:-1] + (7, )) + + _, pre = _get_p(elc_grid, elc_values, gas_gamma=gas_gamma, num_moms=num_moms) + _, kee = _get_ke(elc_grid, elc_values, gas_gamma=gas_gamma, num_moms=num_moms) + _, pri = _get_p(ion_grid, ion_values, gas_gamma=gas_gamma, num_moms=num_moms) + _, kei = _get_ke(ion_grid, ion_values, gas_gamma=gas_gamma, num_moms=num_moms) + _, esq = numerics.mag_sq(field_grid, field_values, coords="0:3") + _, bsq = numerics.mag_sq(field_grid, field_values, coords="3:6") + + out[..., 0] = np.squeeze(pre) + out[..., 1] = np.squeeze(kee) + out[..., 2] = np.squeeze(pri) + out[..., 3] = np.squeeze(kei) + out[..., 4] = np.squeeze(esq / 2.0) + out[..., 5] = np.squeeze(bsq / 2.0) + out[..., 6] = np.squeeze(pre + kee + pri + kei + esq / 2.0 + bsq / 2.0) + + return list(field_grid), out + + +def _accumulate_current( + grid: list[np.ndarray], + values: np.ndarray, + *, + qbym: bool = False, + charge: float | None = None, + mass: float | None = None, +) -> tuple[list[np.ndarray], np.ndarray]: + """Scale a species' moment data into its contribution to the current. + + Args: + grid: Species moment grid. + values: Species moment array. + qbym: If ``True``, scale by the charge-to-mass ratio ``charge / mass`` + (appropriate for fluid moment data, which already carries a mass + factor in the density); otherwise scale by ``-1.0``. + charge: Particle charge, required when ``qbym`` is ``True``. + mass: Particle mass, required (and must be nonzero) when ``qbym`` is + ``True``. + + Returns: + ``(grid, values)`` holding the current contribution. + """ + if qbym and mass and charge is not None: + factor = charge / mass + else: + factor = -1.0 + + return list(grid), factor * values + + +# ---------------------------------------------------------------- GData verbs +def energetics(elc: "GDataState", + ion: "GDataState", + field: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + num_moms: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Decompose energy (kinetic, thermal, EM) for a two-species plasma. + + Splits the plasma energy into its constituent parts for a two-species + (electron/ion) plasma plus an EM field. The result carries the EM + field's grid and metadata and has seven components, in order: + + 0. electron thermal energy + 1. electron kinetic energy + 2. ion thermal energy + 3. ion kinetic energy + 4. electric field energy (``|E|^2 / 2``) + 5. magnetic field energy (``|B|^2 / 2``) + 6. total energy (sum of the above) + + Args: + elc: Electron fluid moments (used to compute thermal pressure and + kinetic energy); must be NumPy-backed. + ion: Ion fluid moments (used to compute thermal pressure and kinetic + energy); must be NumPy-backed. + field: EM field whose components 0:3 are the electric field and 3:6 + are the magnetic field; its grid/metadata are carried to the output. + Must be NumPy-backed. + gas_gamma: Adiabatic index, forwarded to the pressure/kinetic-energy + calculation for both species. + num_moms: Number of moments (5 or 10) for both species; inferred from + the component count when ``None``. + inplace: mutate and return ``field`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A seven-component dataset of the energy decomposition. + + Raises: + ValueError: if any input is native modal (gkyl-backed). + """ + _require_field_domain(elc, "energetics", _REASON) + _require_field_domain(ion, "energetics", _REASON) + _require_field_domain(field, "energetics", _REASON) + grid, values = _energetics(elc.grid, + elc.values, + ion.grid, + ion.values, + field.grid, + field.values, + gas_gamma=gas_gamma, + num_moms=num_moms) + return field._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def accumulate_current(data: "GDataState", + *, + qbym: bool = False, + charge: float | None = None, + mass: float | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Accumulate current from species moments. + + Scales the species' momentum/flow moments by a per-species factor to + form its contribution to the current. By default the factor is ``-1.0``; + with ``qbym=True`` (and ``charge``/``mass`` given) the charge/mass ratio + is used instead. Should be used with ``qbym=True`` for fluid data. + + Args: + data: A species dataset carrying the flow/momentum moments to scale; + must be NumPy-backed. + qbym: When True, scale by the charge-to-mass ratio (q/m); otherwise + scale by ``-1.0``. Set True for fluid data. + charge: Particle charge, required when ``qbym`` is True. + mass: Particle mass, required (and must be nonzero) when ``qbym`` is + True. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset of the scaled current contribution. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed); if ``qbym`` is + True and ``charge``/``mass`` are not both given (a nonzero ``mass``). + """ + if data.backend == "gkyl": + raise ValueError( + "accumulate_current operates on interpolated (NumPy) values; call " + ".interpolate() first -- scaling raw DG coefficients by a per-species " + "factor is still valid numerically, but this verb is field-domain " + "only.") + if qbym and (charge is None or not mass): + raise ValueError( + "accumulate_current: qbym=True requires both 'charge' and a " + f"nonzero 'mass' -- got charge={charge!r}, mass={mass!r}.") + grid, values = _accumulate_current(data.grid, + data.values, + qbym=qbym, + charge=charge, + mass=mass) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/diagnostics/mom/plasma.py b/src/postgkyl/diagnostics/mom/plasma.py new file mode 100644 index 00000000..03a721b2 --- /dev/null +++ b/src/postgkyl/diagnostics/mom/plasma.py @@ -0,0 +1,616 @@ +"""Plasma parameters: field magnitude, thermal/Alfven velocity, cyclotron and +plasma frequency, inertial length, Debye length, gyroradius, plasma beta. + +These never had a verb layer of their own (only the array math lived in the +old ``models`` package) -- every public function here is a fresh GData-facing +wrapper (species/field datasets in, ``GDataState`` out) over that moved +array math. + +The old ``postgkeyll.tools.params`` functions read ``mass``/``charge``/ +``mu_0``/``epsilon_0`` from a ``GData.ctx`` dict, falling back to a keyword +argument when the context held nothing. These are pure keyword-only +arguments instead -- no ctx, no fallback chain. A consequence of dropping the +GData/ctx duality is that a few old parameters were never anything but ctx +lookups (unused otherwise) and are dropped here because keeping them would +misstate what the function actually needs (Doctrine IV): ``omegaC`` does not +take ``species`` (only ``field`` values were ever used), ``omegaP``/``d``/ +``lambdaD`` do not take ``field`` (only ``species`` values were ever used), +and ``rho`` drops the never-referenced ``epsilon_0`` parameter. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ... import numerics +from ...gdatastate.guards import require_field_domain as _require_field_domain +from .five_moment import _get_density, _get_temp +from .mhd import _get_mhd_temp + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = "computing plasma parameters from raw DG coefficients would mix basis functions" + + +# --------------------------------------------------------- array-level math +def _get_magB(field_grid: list[np.ndarray], + field_values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the magnitude of the magnetic field ``|B|``. + + Args: + field_grid: EM field grid. + field_values: EM field array laid out ``[Ex, Ey, Ez, Bx, By, Bz, ...]``; + components 3:6 are used. + + Returns: + ``(grid, values)`` holding ``|B| = sqrt(Bx**2 + By**2 + Bz**2)``. + """ + b_values = field_values[..., 3:6] + _, mag_B_sq = numerics.mag_sq(field_grid, b_values) + return list(field_grid), np.sqrt(mag_B_sq) + + +def _get_vt(species_grid: list[np.ndarray], + species_values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, + mhd: bool = False) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the thermal velocity ``v_th = sqrt(2 T/m)`` (or ``sqrt(T/m)`` + when ``no_sqrt2`` is ``True``) of a species. + + Args: + species_grid: Species moment grid. + species_values: Species moment array (5- or 10-moment, or MHD when + ``mhd=True``). + gas_gamma: Adiabatic index used when computing the temperature/pressure. + num_moms: Number of moments (5 or 10); inferred when ``None``. + mass: Particle mass. + mu_0: Vacuum permeability, forwarded to the MHD temperature when + ``mhd=True``. + no_sqrt2: Omit the conventional ``sqrt(2)`` scale factor. + mhd: If ``True``, compute the temperature from MHD moments; otherwise + use the fluid moments. + + Returns: + ``(grid, values)`` holding the thermal velocity field. + """ + if mhd: + out_grid, temp = _get_mhd_temp(species_grid, + species_values, + gas_gamma=gas_gamma, + mu_0=mu_0) + else: + out_grid, temp = _get_temp(species_grid, + species_values, + gas_gamma=gas_gamma, + num_moms=num_moms) + + out_values = np.sqrt(temp / mass) + if not no_sqrt2: + out_values = out_values * np.sqrt(2.0) + + return out_grid, out_values + + +def _get_vA(species_grid: list[np.ndarray], + species_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, + *, + mu_0: float = 1.0) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the Alfven velocity ``v_A = |B| / sqrt(mu_0 * rho)``. + + Fluid moment data already includes the mass factor in the density. + """ + _, magB = _get_magB(field_grid, field_values) + out_grid, rho = _get_density(species_grid, species_values) + return out_grid, magB / np.sqrt(mu_0 * rho) + + +def _get_omegaC( + field_grid: list[np.ndarray], + field_values: np.ndarray, + *, + mass: float = 1.0, + charge: float = 1.0, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the cyclotron (gyro) frequency ``omega_c = |q| * |B| / m``.""" + out_grid, magB = _get_magB(field_grid, field_values) + return out_grid, abs(charge) * magB / mass + + +def _get_omegaP( + species_grid: list[np.ndarray], + species_values: np.ndarray, + *, + mass: float = 1.0, + charge: float = 1.0, + epsilon_0: float = 1.0, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the plasma frequency + ``omega_p = sqrt(q**2 * n / (m**2 * epsilon_0))``. + + Fluid moment data already includes the mass factor in the density. + """ + out_grid, rho = _get_density(species_grid, species_values) + qbym2 = charge**2 / mass**2 + return out_grid, np.sqrt(qbym2 * rho / epsilon_0) + + +def _get_d(species_grid: list[np.ndarray], + species_values: np.ndarray, + *, + mass: float = 1.0, + charge: float = 1.0, + epsilon_0: float = 1.0, + mu_0: float = 1.0) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the inertial (skin-depth) length ``d = c / omega_p``, with + ``c = 1 / sqrt(epsilon_0 * mu_0)``.""" + out_grid, omegaP = _get_omegaP(species_grid, + species_values, + mass=mass, + charge=charge, + epsilon_0=epsilon_0) + light_speed = 1.0 / np.sqrt(epsilon_0 * mu_0) + return out_grid, light_speed / omegaP + + +def _get_lambdaD( + species_grid: list[np.ndarray], + species_values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + charge: float = 1.0, + epsilon_0: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the Debye length ``lambda_D = v_th / omega_p``. + + Unless ``no_sqrt2`` is set, the extra ``sqrt(2)`` factor carried by + ``v_th`` is divided back out, so the conventional Debye length is + returned. + """ + _, omegaP = _get_omegaP(species_grid, + species_values, + mass=mass, + charge=charge, + epsilon_0=epsilon_0) + out_grid, vt = _get_vt(species_grid, + species_values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + mu_0=mu_0, + no_sqrt2=no_sqrt2) + out_values = vt / omegaP + if not no_sqrt2: + out_values = out_values / np.sqrt(2.0) + + return out_grid, out_values + + +def _get_rho(species_grid: list[np.ndarray], + species_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + charge: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the gyroradius (Larmor radius) ``rho = v_th / omega_c``. + + When ``no_sqrt2`` is set the result is multiplied by ``sqrt(2)`` so the + gyroradius stays consistent with a ``sqrt(2)``-scaled thermal velocity. + """ + _, omegaC = _get_omegaC(field_grid, field_values, mass=mass, charge=charge) + out_grid, vt = _get_vt(species_grid, + species_values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + mu_0=mu_0, + no_sqrt2=no_sqrt2) + + out_values = vt / omegaC + if no_sqrt2: + out_values = out_values * np.sqrt(2.0) + + return out_grid, out_values + + +def _get_beta( + species_grid: list[np.ndarray], + species_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the plasma beta ``v_th**2 / v_A**2``. + + When ``no_sqrt2`` is set the result is multiplied by ``2`` to account + for the missing ``sqrt(2)`` factor in the thermal velocity. + """ + _, v_A = _get_vA(species_grid, + species_values, + field_grid, + field_values, + mu_0=mu_0) + out_grid, vt = _get_vt(species_grid, + species_values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + mu_0=mu_0, + no_sqrt2=no_sqrt2) + out_values = vt**2 / v_A**2 + if no_sqrt2: + out_values = out_values * 2.0 + + return out_grid, out_values + + +# ---------------------------------------------------------------- GData verbs +def magB(field: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Magnitude of the magnetic field ``|B|``. + + Args: + field: EM field data (components 3:6 are ``Bx, By, Bz``); must be + NumPy-backed. + inplace: mutate and return ``field`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of ``|B|``. + + Raises: + ValueError: if ``field`` is native modal (gkyl-backed). + """ + _require_field_domain(field, "magB", _REASON) + grid, values = _get_magB(field.grid, field.values) + return field._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def vt(species: "GDataState", + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, + mhd: bool = False, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Thermal velocity ``v_th = sqrt(2 T/m)`` of a species. + + Args: + species: Species moment data (5- or 10-moment, or MHD when ``mhd=True``); + must be NumPy-backed. + gas_gamma: Adiabatic index used when computing the temperature/pressure. + num_moms: Number of moments (5 or 10); inferred when ``None``. + mass: Particle mass. + mu_0: Vacuum permeability, forwarded to the MHD temperature when + ``mhd=True``. + no_sqrt2: Omit the conventional ``sqrt(2)`` scale factor. + mhd: If ``True``, compute the temperature from MHD moments; otherwise + use the fluid moments. + inplace: mutate and return ``species`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the thermal velocity. + + Raises: + ValueError: if ``species`` is native modal (gkyl-backed). + """ + _require_field_domain(species, "vt", _REASON) + grid, values = _get_vt(species.grid, + species.values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + mu_0=mu_0, + no_sqrt2=no_sqrt2, + mhd=mhd) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def vA(species: "GDataState", + field: "GDataState", + *, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Alfven velocity ``v_A = |B| / sqrt(mu_0 * rho)``. + + Args: + species: Species moment data providing the density; must be + NumPy-backed. + field: EM field data providing ``|B|``; must be NumPy-backed. + mu_0: Vacuum permeability. + inplace: mutate and return ``species`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the Alfven velocity. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(species, "vA", _REASON) + _require_field_domain(field, "vA", _REASON) + grid, values = _get_vA(species.grid, + species.values, + field.grid, + field.values, + mu_0=mu_0) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def omegaC(field: "GDataState", + *, + mass: float = 1.0, + charge: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Cyclotron (gyro) frequency ``omega_c = |q| * |B| / m``. + + Args: + field: EM field data providing ``Bx, By, Bz``; must be NumPy-backed. + mass: Particle mass. + charge: Particle charge; only its magnitude affects the result. + inplace: Mutate and return ``field`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the cyclotron frequency. + + Raises: + ValueError: if ``field`` is native modal (gkyl-backed). + """ + _require_field_domain(field, "omegaC", _REASON) + grid, values = _get_omegaC(field.grid, field.values, mass=mass, charge=charge) + return field._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def omegaP(species: "GDataState", + *, + mass: float = 1.0, + charge: float = 1.0, + epsilon_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Plasma frequency ``omega_p = sqrt(q**2 * n / (m**2 * epsilon_0))``. + + Args: + species: Fluid moment data providing mass density; must be NumPy-backed. + mass: Particle mass. + charge: Particle charge. + epsilon_0: Vacuum permittivity. + inplace: Mutate and return ``species`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the plasma frequency. + + Raises: + ValueError: if ``species`` is native modal (gkyl-backed). + """ + _require_field_domain(species, "omegaP", _REASON) + grid, values = _get_omegaP(species.grid, + species.values, + mass=mass, + charge=charge, + epsilon_0=epsilon_0) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def d(species: "GDataState", + *, + mass: float = 1.0, + charge: float = 1.0, + epsilon_0: float = 1.0, + mu_0: float = 1.0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Inertial (skin-depth) length ``d = c / omega_p``. + + Args: + species: Fluid moment data providing mass density; must be NumPy-backed. + mass: Particle mass. + charge: Particle charge. + epsilon_0: Vacuum permittivity. + mu_0: Vacuum permeability used in ``c = 1/sqrt(epsilon_0*mu_0)``. + inplace: Mutate and return ``species`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the inertial length. + + Raises: + ValueError: if ``species`` is native modal (gkyl-backed). + """ + _require_field_domain(species, "d", _REASON) + grid, values = _get_d(species.grid, + species.values, + mass=mass, + charge=charge, + epsilon_0=epsilon_0, + mu_0=mu_0) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def lambdaD(species: "GDataState", + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + charge: float = 1.0, + epsilon_0: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Debye length ``lambda_D = v_th / omega_p``. + + Args: + species: Fluid moment data; must be NumPy-backed. + gas_gamma: Adiabatic index used to compute thermal velocity. + num_moms: Number of fluid moments (5 or 10); inferred when ``None``. + mass: Particle mass. + charge: Particle charge. + epsilon_0: Vacuum permittivity used in the plasma frequency. + mu_0: Vacuum permeability forwarded to thermal-velocity calculation. + no_sqrt2: Use ``sqrt(T/m)`` internally; the returned conventional Debye + length remains unchanged. + inplace: Mutate and return ``species`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the Debye length. + + Raises: + ValueError: if ``species`` is native modal (gkyl-backed). + """ + _require_field_domain(species, "lambdaD", _REASON) + grid, values = _get_lambdaD(species.grid, + species.values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + charge=charge, + epsilon_0=epsilon_0, + mu_0=mu_0, + no_sqrt2=no_sqrt2) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def rho(species: "GDataState", + field: "GDataState", + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + charge: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Gyroradius (Larmor radius) ``rho = v_th / omega_c``. + + Args: + species: Fluid moment data used for thermal velocity; must be + NumPy-backed. + field: EM field data used for cyclotron frequency; must be NumPy-backed. + gas_gamma: Adiabatic index used to compute thermal velocity. + num_moms: Number of fluid moments (5 or 10); inferred when ``None``. + mass: Particle mass. + charge: Particle charge. + mu_0: Vacuum permeability forwarded to thermal-velocity calculation. + no_sqrt2: Use ``sqrt(T/m)`` internally; the result is normalized to the + ``sqrt(2*T/m)`` convention either way. + inplace: Mutate and return ``species`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the gyroradius. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(species, "rho", _REASON) + _require_field_domain(field, "rho", _REASON) + grid, values = _get_rho(species.grid, + species.values, + field.grid, + field.values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + charge=charge, + mu_0=mu_0, + no_sqrt2=no_sqrt2) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def beta(species: "GDataState", + field: "GDataState", + *, + gas_gamma: float = 5.0 / 3.0, + num_moms: int | None = None, + mass: float = 1.0, + mu_0: float = 1.0, + no_sqrt2: bool = False, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Plasma beta ``v_th**2 / v_A**2``. + + Args: + species: Fluid moment data used for thermal velocity and density; must + be NumPy-backed. + field: EM field data used for Alfven velocity; must be NumPy-backed. + gas_gamma: Adiabatic index used to compute thermal velocity. + num_moms: Number of fluid moments (5 or 10); inferred when ``None``. + mass: Particle mass. + mu_0: Vacuum permeability. + no_sqrt2: Use ``sqrt(T/m)`` internally; the result is normalized to the + ``sqrt(2*T/m)`` convention either way. + inplace: Mutate and return ``species`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of plasma beta. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(species, "beta", _REASON) + _require_field_domain(field, "beta", _REASON) + grid, values = _get_beta(species.grid, + species.values, + field.grid, + field.values, + gas_gamma=gas_gamma, + num_moms=num_moms, + mass=mass, + mu_0=mu_0, + no_sqrt2=no_sqrt2) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/diagnostics/mom/rotations.py b/src/postgkyl/diagnostics/mom/rotations.py new file mode 100644 index 00000000..0ecd205d --- /dev/null +++ b/src/postgkyl/diagnostics/mom/rotations.py @@ -0,0 +1,222 @@ +"""Vector rotation parallel/perpendicular to a reference (e.g. the magnetic +field). + +For a field ``u`` and a rotator ``v`` (assumed three-component, last axis), +``parrotate`` computes the projection of ``u`` onto ``v``'s direction, +``(u . v_hat) v_hat``; ``perprotate`` is the remainder, ``u - (u . v_hat) +v_hat``. + +Note: :mod:`postgkyl.numerics.rotation_matrix` builds a matrix whose first +row is the *elementwise sign* of its input, not a true unit vector (see its +own tests) -- using it here would change the projection's numerical result, +so this module keeps the original dot-product formula instead (Doctrine: +copy numerics verbatim). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ...gdatastate.guards import require_field_domain as _require_field_domain + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = "rotating raw DG coefficients would mix basis functions" + + +# --------------------------------------------------------- array-level math +def _parrotate( + grid: list[np.ndarray], + values: np.ndarray, + rotator_values: np.ndarray, + *, + rotate_coords: str = "0:3", +) -> tuple[list[np.ndarray], np.ndarray]: + """Rotate a three-component field into the direction of a rotator field. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension. + values: Three-component field to rotate (last axis is components). + rotator_values: Field providing the rotation direction, on the same + grid as ``values``. + rotate_coords: ``"start:end"`` slice of ``rotator_values``'s component + axis to use as the rotation direction (e.g. ``"3:6"`` to rotate into + a magnetic field stored after three electric-field components). + + Returns: + ``(grid, values)`` holding the parallel component + ``(u . v_hat) v_hat``. + + Raises: + ValueError: If ``values`` or the sliced ``rotator_values`` do not have + exactly three components. + """ + lo, hi = rotate_coords.split(":") + valuesrot = rotator_values[..., slice(int(lo), int(hi))] + + if values.shape[-1] != 3 or valuesrot.shape[-1] != 3: + raise ValueError( + "parrotate requires three-component vector fields; data has " + f"{values.shape[-1]:d} components, rotator (after 'rotate_coords' " + f"slicing) has {valuesrot.shape[-1]:d}") + + scale = np.sum(values * valuesrot, axis=-1) / np.sum(valuesrot * valuesrot, + axis=-1) + outrot = scale[..., np.newaxis] * valuesrot + + return list(grid), outrot + + +def _perprotate( + grid: list[np.ndarray], + values: np.ndarray, + rotator_values: np.ndarray, + *, + rotate_coords: str = "0:3", +) -> tuple[list[np.ndarray], np.ndarray]: + """Rotate a three-component field perpendicular to a rotator field. + + Computed as the remainder after :func:`_parrotate`: + ``u - (u . v_hat) v_hat``. + """ + grid, par = _parrotate(grid, + values, + rotator_values, + rotate_coords=rotate_coords) + return grid, values - par + + +# ---------------------------------------------------------------- GData verbs +def parrotate(array: "GDataState", + rotator: "GDataState", + *, + coords: str = "0:3", + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Component of ``array`` parallel to ``rotator``: ``(u . v_hat) v_hat``. + + Projects the three-component vector field ``array`` (u) onto the unit + vector of the ``rotator`` field (v), returning the parallel vector + ``(u . v_hat) v_hat`` with its x, y, z components. Both fields are + assumed to be three-component with components on the last axis. + + Args: + array: The three-component vector field to be rotated/projected; must + be NumPy-backed. + rotator: The field defining the rotation direction; must be + NumPy-backed. + coords: Half-open 'lo:hi' slice string selecting which ``rotator`` + components form the direction vector. Defaults to '0:3'; use '3:6' + to rotate along the magnetic field of a six-component EM field. + inplace: mutate and return ``array`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A three-component dataset of the parallel projection. + + Raises: + ValueError: if either input is native modal (gkyl-backed), or the + component counts do not match a three-component field. + """ + _require_field_domain(array, "parrotate", _REASON) + _require_field_domain(rotator, "parrotate", _REASON) + grid, values = _parrotate(array.grid, + array.values, + rotator.values, + rotate_coords=coords) + return array._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def perprotate(array: "GDataState", + rotator: "GDataState", + *, + coords: str = "0:3", + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Component of ``array`` perpendicular to ``rotator``: + ``u - (u . v_hat) v_hat``. + + Both fields are assumed to be three-component with components on the + last axis. + + Args: + array: The three-component vector field to be rotated/projected; must + be NumPy-backed. + rotator: The field defining the rotation direction; must be + NumPy-backed. + coords: Half-open 'lo:hi' slice string selecting which ``rotator`` + components form the direction vector. Defaults to '0:3'; use '3:6' + to rotate along the magnetic field of a six-component EM field. + inplace: mutate and return ``array`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A three-component dataset of the perpendicular component. + + Raises: + ValueError: if either input is native modal (gkyl-backed), or the + component counts do not match a three-component field. + """ + _require_field_domain(array, "perprotate", _REASON) + _require_field_domain(rotator, "perprotate", _REASON) + grid, values = _perprotate(array.grid, + array.values, + rotator.values, + rotate_coords=coords) + return array._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def bparrotate(array: "GDataState", + field: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Project an array parallel to the magnetic field. + + Args: + array: Vector or tensor dataset to project. + field: Electromagnetic field whose components 3 through 5 are magnetic. + inplace: Mutate and return ``array`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + return parrotate(array, + field, + coords="3:6", + inplace=inplace, + tag=tag, + label=label) + + +def bperprotate(array: "GDataState", + field: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Project an array perpendicular to the magnetic field. + + Args: + array: Vector or tensor dataset to project. + field: Electromagnetic field whose components 3 through 5 are magnetic. + inplace: Mutate and return ``array`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + return perprotate(array, + field, + coords="3:6", + inplace=inplace, + tag=tag, + label=label) + + +__all__ = ["parrotate", "perprotate", "bparrotate", "bperprotate"] diff --git a/src/postgkyl/diagnostics/mom/ten_moment.py b/src/postgkyl/diagnostics/mom/ten_moment.py new file mode 100644 index 00000000..527a572e --- /dev/null +++ b/src/postgkyl/diagnostics/mom/ten_moment.py @@ -0,0 +1,770 @@ +"""Ten-moment diagnostics -- the five-moment set (fixed to 10-moment data) +plus the pressure tensor, field-aligned pressure, and agyrotropy. + +10-moment fluid data is laid out ``[rho, mx, my, mz, Pxx, Pxy, Pxz, Pyy, Pyz, +Pzz]``; the pressure tensor components below subtract the bulk-flow (ram) +contribution from the raw second moments. ``p_par``/``p_perp``/``agyro`` +then take an already-built 6-component pressure tensor (``P_xx, P_xy, P_xz, +P_yy, P_yz, P_zz``) and a 3-component magnetic field. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ... import numerics +from ...gdatastate.guards import require_field_domain as _require_field_domain +from .five_moment import ( + _get_density, + _get_vx, + _get_vy, + _get_vz, + _get_p, + _get_ke, + _get_temp, + _get_sound, + _get_mach, + density, + xvel, + yvel, + zvel, + vel, +) + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = ("extracting primitive variables from raw DG coefficients would " + "mix basis functions") +_AGYRO_REASON = ("computing agyrotropy from raw DG coefficients would mix " + "basis functions") + + +# --------------------------------------------------------- array-level math +def _get_pxx(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """``P_xx = M_xx - rho * vx * vx`` (component 4 of 10-moment data).""" + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + return list(grid), values[..., 4, np.newaxis] - rho * vx * vx + + +def _get_pxy(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """``P_xy = M_xy - rho * vx * vy`` (component 5 of 10-moment data).""" + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + _, vy = _get_vy(grid, values) + return list(grid), values[..., 5, np.newaxis] - rho * vx * vy + + +def _get_pxz(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """``P_xz = M_xz - rho * vx * vz`` (component 6 of 10-moment data).""" + _, rho = _get_density(grid, values) + _, vx = _get_vx(grid, values) + _, vz = _get_vz(grid, values) + return list(grid), values[..., 6, np.newaxis] - rho * vx * vz + + +def _get_pyy(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """``P_yy = M_yy - rho * vy * vy`` (component 7 of 10-moment data).""" + _, rho = _get_density(grid, values) + _, vy = _get_vy(grid, values) + return list(grid), values[..., 7, np.newaxis] - rho * vy * vy + + +def _get_pyz(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """``P_yz = M_yz - rho * vy * vz`` (component 8 of 10-moment data).""" + _, rho = _get_density(grid, values) + _, vy = _get_vy(grid, values) + _, vz = _get_vz(grid, values) + return list(grid), values[..., 8, np.newaxis] - rho * vy * vz + + +def _get_pzz(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """``P_zz = M_zz - rho * vz * vz`` (component 9 of 10-moment data).""" + _, rho = _get_density(grid, values) + _, vz = _get_vz(grid, values) + return list(grid), values[..., 9, np.newaxis] - rho * vz * vz + + +def _get_pij(grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Full symmetric pressure tensor, packed + ``(P_xx, P_xy, P_xz, P_yy, P_yz, P_zz)``.""" + out_values = np.zeros(values[..., 4:10].shape) + _, pxx = _get_pxx(grid, values) + _, pxy = _get_pxy(grid, values) + _, pxz = _get_pxz(grid, values) + _, pyy = _get_pyy(grid, values) + _, pyz = _get_pyz(grid, values) + _, pzz = _get_pzz(grid, values) + + out_values[..., 0] = np.squeeze(pxx) + out_values[..., 1] = np.squeeze(pxy) + out_values[..., 2] = np.squeeze(pxz) + out_values[..., 3] = np.squeeze(pyy) + out_values[..., 4] = np.squeeze(pyz) + out_values[..., 5] = np.squeeze(pzz) + + return list(grid), out_values + + +def _get_p_par( + p_grid: list[np.ndarray], + p_values: np.ndarray, + b_grid: list[np.ndarray], + b_values: np.ndarray, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the pressure parallel to the magnetic field. + + Projects the pressure tensor onto the magnetic-field direction: + ``p_par = (b . P . b) / |B|**2``. + + Args: + p_grid: Pressure-tensor grid. + p_values: 6-component pressure tensor + ``(P_xx, P_xy, P_xz, P_yy, P_yz, P_zz)``. + b_grid: Magnetic-field grid. + b_values: 3-component magnetic field ``(Bx, By, Bz)``. + + Returns: + ``(grid, values)`` holding the parallel pressure field. + """ + p_xx = p_values[..., 0, np.newaxis] + p_xy = p_values[..., 1, np.newaxis] + p_xz = p_values[..., 2, np.newaxis] + p_yy = p_values[..., 3, np.newaxis] + p_yz = p_values[..., 4, np.newaxis] + p_zz = p_values[..., 5, np.newaxis] + + b_x = b_values[..., 0, np.newaxis] + b_y = b_values[..., 1, np.newaxis] + b_z = b_values[..., 2, np.newaxis] + + grid, mag_b_sq = numerics.mag_sq(b_grid, b_values) + + out = (b_x * b_x * p_xx + b_y * b_y * p_yy + b_z * b_z * p_zz + 2.0 * + (b_x * b_y * p_xy + b_x * b_z * p_xz + b_y * b_z * p_yz)) / mag_b_sq + return grid, out + + +def _get_gkyl_10m_p_par( + species_grid: list[np.ndarray], + species_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the parallel pressure directly from raw 10-moment species and + EM field data (whose components 3:6 are ``(Bx, By, Bz)``).""" + p_grid, p_values = _get_pij(species_grid, species_values) + b_values = field_values[..., 3:6] + return _get_p_par(p_grid, p_values, field_grid, b_values) + + +def _get_p_perp( + p_grid: list[np.ndarray], + p_values: np.ndarray, + b_grid: list[np.ndarray], + b_values: np.ndarray, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the pressure perpendicular to the magnetic field. + + Uses the trace of the pressure tensor and the parallel pressure: + ``p_perp = (P_xx + P_yy + P_zz - p_par) / 2``. + """ + p_xx = p_values[..., 0, np.newaxis] + p_yy = p_values[..., 3, np.newaxis] + p_zz = p_values[..., 5, np.newaxis] + + grid, p_par = _get_p_par(p_grid, p_values, b_grid, b_values) + + return grid, (p_xx + p_yy + p_zz - p_par) / 2.0 + + +def _get_gkyl_10m_p_perp( + species_grid: list[np.ndarray], + species_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compute the perpendicular pressure directly from raw 10-moment species + and EM field data (whose components 3:6 are ``(Bx, By, Bz)``).""" + p_grid, p_values = _get_pij(species_grid, species_values) + b_values = field_values[..., 3:6] + return _get_p_perp(p_grid, p_values, field_grid, b_values) + + +def _get_agyro(p_grid: list[np.ndarray], + p_values: np.ndarray, + b_grid: list[np.ndarray], + b_values: np.ndarray, + *, + measure: str = "swisdak") -> tuple[list[np.ndarray], np.ndarray]: + """Compute the agyrotropy of the pressure tensor. + + The ``'swisdak'`` measure uses the tensor invariants and parallel pressure + as in Appendix A of Swisdak (2015). The ``'frobenius'`` measure is the + Frobenius norm of the non-gyrotropic part of the pressure tensor, + normalized by the gyrotropic part. + + Args: + p_grid: Pressure-tensor grid. + p_values: 6-component pressure tensor + ``(P_xx, P_xy, P_xz, P_yy, P_yz, P_zz)``. + b_grid: Magnetic-field grid. + b_values: 3-component magnetic field ``(Bx, By, Bz)``. + measure: ``'swisdak'`` (default) or ``'frobenius'`` (case-insensitive). + + Returns: + ``(grid, values)`` holding the agyrotropy field. + + Raises: + ValueError: If ``measure`` is neither ``'swisdak'`` nor ``'frobenius'``. + """ + p_xx = p_values[..., 0, np.newaxis] + p_xy = p_values[..., 1, np.newaxis] + p_xz = p_values[..., 2, np.newaxis] + p_yy = p_values[..., 3, np.newaxis] + p_yz = p_values[..., 4, np.newaxis] + p_zz = p_values[..., 5, np.newaxis] + + b_x = b_values[..., 0, np.newaxis] + b_y = b_values[..., 1, np.newaxis] + b_z = b_values[..., 2, np.newaxis] + + grid, mag_b_sq = numerics.mag_sq(b_grid, b_values) + _, p_par = _get_p_par(p_grid, p_values, b_grid, b_values) + _, p_perp = _get_p_perp(p_grid, p_values, b_grid, b_values) + + measure_lower = measure.lower() + if measure_lower == "swisdak": + I1 = p_xx + p_yy + p_zz + I2 = (p_xx * p_yy + p_xx * p_zz + p_yy * p_zz - + (p_xy * p_xy + p_xz * p_xz + p_yz * p_yz)) + # Tensor algebra of Appendix A of Swisdak 2015. + out = np.sqrt(1 - 4 * I2 / ((I1 - p_par) * (I1 + 3 * p_par))) + elif measure_lower == "frobenius": + p_ixx = p_xx - (p_par * b_x * b_x / mag_b_sq + p_perp * + (1 - b_x * b_x / mag_b_sq)) + p_ixy = p_xy - (p_par * b_x * b_y / mag_b_sq + p_perp * + (0 - b_x * b_y / mag_b_sq)) + p_ixz = p_xz - (p_par * b_x * b_z / mag_b_sq + p_perp * + (0 - b_x * b_z / mag_b_sq)) + p_iyy = p_yy - (p_par * b_y * b_y / mag_b_sq + p_perp * + (1 - b_y * b_y / mag_b_sq)) + p_iyz = p_yz - (p_par * b_y * b_z / mag_b_sq + p_perp * + (0 - b_y * b_z / mag_b_sq)) + p_izz = p_zz - (p_par * b_z * b_z / mag_b_sq + p_perp * + (1 - b_z * b_z / mag_b_sq)) + out = (np.sqrt(p_ixx**2 + 2 * p_ixy**2 + 2 * p_ixz**2 + p_iyy**2 + + 2 * p_iyz**2 + p_izz**2) / + np.sqrt(2 * p_perp**2 + 4 * p_par * p_perp)) + else: + raise ValueError( + f"Measure specified is {measure_lower:s}; it needs to be either " + "'swisdak' or 'frobenius'") + + return grid, out + + +def _get_gkyl_10m_agyro( + species_grid: list[np.ndarray], + species_values: np.ndarray, + field_grid: list[np.ndarray], + field_values: np.ndarray, + *, + measure: str = "swisdak") -> tuple[list[np.ndarray], np.ndarray]: + """Compute the agyrotropy directly from raw 10-moment species and EM field + data (whose components 3:6 are ``(Bx, By, Bz)``).""" + p_grid, p_values = _get_pij(species_grid, species_values) + b_values = field_values[..., 3:6] + return _get_agyro(p_grid, p_values, field_grid, b_values, measure=measure) + + +# ---------------------------------------------------------------- GData verbs +def pressure(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Scalar pressure (trace of the pressure tensor over three) from + 10-moment fluid data. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + gas_gamma: Unused compatibility parameter. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the scalar pressure. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pressure", _REASON) + grid, values = _get_p(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=10) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def ke(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Kinetic (bulk-flow) energy density from 10-moment fluid data. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + gas_gamma: Unused compatibility parameter. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the bulk-flow energy density. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "ke", _REASON) + grid, values = _get_ke(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=10) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def temp(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Temperature ``T = p / rho`` from 10-moment fluid data. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + gas_gamma: Unused compatibility parameter. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the temperature. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "temp", _REASON) + grid, values = _get_temp(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=10) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def sound(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Sound speed ``c_s = sqrt(gas_gamma * p / rho)`` from 10-moment data. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + gas_gamma: Adiabatic index used in the sound speed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the sound speed. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "sound", _REASON) + grid, values = _get_sound(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=10) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def mach(data: "GDataState", + *, + gas_gamma: float = 5.0 / 3, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Sonic Mach number ``M = |v| / c_s`` from 10-moment data. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + gas_gamma: Adiabatic index used in the sound speed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of the Mach number. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "mach", _REASON) + grid, values = _get_mach(data.grid, + data.values, + gas_gamma=gas_gamma, + num_moms=10) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pxx(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """``P_xx`` pressure-tensor component. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``P_xx``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pxx", _REASON) + grid, values = _get_pxx(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pxy(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """``P_xy`` pressure-tensor component. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``P_xy``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pxy", _REASON) + grid, values = _get_pxy(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pxz(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """``P_xz`` pressure-tensor component. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``P_xz``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pxz", _REASON) + grid, values = _get_pxz(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pyy(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """``P_yy`` pressure-tensor component. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``P_yy``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pyy", _REASON) + grid, values = _get_pyy(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pyz(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """``P_yz`` pressure-tensor component. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``P_yz``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pyz", _REASON) + grid, values = _get_pyz(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pzz(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """``P_zz`` pressure-tensor component. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A single-component dataset of ``P_zz``. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pzz", _REASON) + grid, values = _get_pzz(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def pressure_tensor(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Full symmetric pressure tensor + ``(P_xx, P_xy, P_xz, P_yy, P_yz, P_zz)``. + + Args: + data: Ten-moment fluid data; must be NumPy-backed. + inplace: Mutate and return ``data`` instead of a new dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A six-component dataset of the symmetric pressure tensor. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + _require_field_domain(data, "pressure_tensor", _REASON) + grid, values = _get_pij(data.grid, data.values) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def p_par(ptensor: "GDataState", + bfield: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Pressure parallel to the magnetic field: ``(b . P . b) / |B|**2``. + + Args: + ptensor: Six-component symmetric pressure tensor (Pxx, Pxy, Pxz, Pyy, + Pyz, Pzz); must be NumPy-backed. + bfield: Magnetic field whose first three components are (Bx, By, Bz); + must be NumPy-backed. + inplace: mutate and return ``ptensor`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the parallel pressure. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(ptensor, "p_par", _REASON) + _require_field_domain(bfield, "p_par", _REASON) + grid, values = _get_p_par(ptensor.grid, ptensor.values, bfield.grid, + bfield.values) + return ptensor._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def p_perp(ptensor: "GDataState", + bfield: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Pressure perpendicular to the magnetic field: + ``(P_xx + P_yy + P_zz - p_par) / 2``. + + Args: + ptensor: Six-component symmetric pressure tensor (Pxx, Pxy, Pxz, Pyy, + Pyz, Pzz); must be NumPy-backed. + bfield: Magnetic field whose first three components are (Bx, By, Bz); + must be NumPy-backed. + inplace: mutate and return ``ptensor`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the perpendicular pressure. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(ptensor, "p_perp", _REASON) + _require_field_domain(bfield, "p_perp", _REASON) + grid, values = _get_p_perp(ptensor.grid, ptensor.values, bfield.grid, + bfield.values) + return ptensor._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def agyro(ptensor: "GDataState", + bfield: "GDataState", + *, + measure: str = "frobenius", + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Agyrotropy from a pressure tensor and an EM field. + + Measures how far the pressure tensor departs from gyrotropy about the + local magnetic field. The field's first three components are used as the + magnetic field direction. + + Args: + ptensor: Six-component symmetric pressure tensor (Pxx, Pxy, Pxz, Pyy, + Pyz, Pzz); must be NumPy-backed. + bfield: Magnetic field whose first three components are (Bx, By, Bz); + must be NumPy-backed. + measure: 'frobenius' (Frobenius norm of the agyrotropic part of the + pressure tensor) or 'swisdak' (the Q measure of Swisdak 2015). + Case-insensitive. + inplace: mutate and return ``ptensor`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the agyrotropy. + + Raises: + ValueError: if either input is native modal (gkyl-backed), or + ``measure`` is not 'frobenius' or 'swisdak'. + """ + _require_field_domain(ptensor, "agyro", _AGYRO_REASON) + _require_field_domain(bfield, "agyro", _AGYRO_REASON) + grid, values = _get_agyro(ptensor.grid, + ptensor.values, + bfield.grid, + bfield.values, + measure=measure) + return ptensor._result(grid, values, inplace=inplace, tag=tag, label=label) + + +def mom_agyro(species: "GDataState", + field: "GDataState", + *, + measure: str = "frobenius", + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Agyrotropy from raw 10-moment species data and an EM field. + + Convenience wrapper that first forms the pressure tensor from raw + 10-moment species data and extracts the magnetic field (components 3:6) + from a Gkeyll EM field, then computes the agyrotropy. + + Args: + species: Raw 10-moment fluid data for a single species (density, + momentum, and the six pressure-tensor moments); must be NumPy-backed. + field: Gkeyll EM field whose components 3:6 are the magnetic field (Bx, + By, Bz); must be NumPy-backed. + measure: 'frobenius' (Frobenius norm of the agyrotropic part of the + pressure tensor) or 'swisdak' (the Q measure of Swisdak 2015). + Case-insensitive. + inplace: mutate and return ``species`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the agyrotropy. + + Raises: + ValueError: if either input is native modal (gkyl-backed), or + ``measure`` is not 'frobenius' or 'swisdak'. + """ + _require_field_domain(species, "mom_agyro", _AGYRO_REASON) + _require_field_domain(field, "mom_agyro", _AGYRO_REASON) + grid, values = _get_gkyl_10m_agyro(species.grid, + species.values, + field.grid, + field.values, + measure=measure) + return species._result(grid, values, inplace=inplace, tag=tag, label=label) + + +VARIABLES = { + "density": density, + "xvel": xvel, + "yvel": yvel, + "zvel": zvel, + "vel": vel, + "pressure": pressure, + "ke": ke, + "temp": temp, + "sound": sound, + "mach": mach, + "pressureTensor": pressure_tensor, + "pxx": pxx, + "pxy": pxy, + "pxz": pxz, + "pyy": pyy, + "pyz": pyz, + "pzz": pzz, +} diff --git a/src/postgkyl/diagnostics/pkpm/__init__.py b/src/postgkyl/diagnostics/pkpm/__init__.py new file mode 100644 index 00000000..be895e5d --- /dev/null +++ b/src/postgkyl/diagnostics/pkpm/__init__.py @@ -0,0 +1,180 @@ +"""PKPM diagnostics -- distribution-function reconstruction from Laguerre +moments. + +Composes the full distribution function ``f(x, v_par, v_perp)`` out of the +Laguerre expansion coefficients ``F0(x, v_par)``, ``F1(x, v_par)`` (hardcoded +for ``l=0``, ``n=0,1``) and the PKPM ``T/m`` moment. See Jimmy Juno's slides: +https://drive.google.com/file/d/1548tLF9o7vyW3bkrsq6FvAMV-8XJvKtY/view + +``load_pkpm`` is the equation-internal loader for PKPM output files (ported +from ``src_bak/postgkyl/loaders/pkpm.py``): it loads the distribution and its +companion ``pkpm_vars`` file, interpolates them, and applies +``laguerre_compose`` + ``kinetic.transform_frame``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ...gdata.gdata import GData +from ...gdatastate.guards import require_field_domain as _require_field_domain +from ..vm.kinetic import transform_frame + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = "composing raw DG coefficients would mix basis functions" + + +# --------------------------------------------------------- array-level math +def _laguerre_compose( + f_grid: list[np.ndarray], + f_values: np.ndarray, + t_over_m_values: np.ndarray, +) -> tuple[list[np.ndarray], np.ndarray]: + """Compose PKPM expansion coefficients into a single distribution function. + + Args: + f_grid: ``[x, v_par]`` nodal coordinate arrays. + f_values: 2-component Laguerre expansion coefficients ``(F0, G)``. + t_over_m_values: PKPM ``T / m`` moment, single component. + + Returns: + ``([x, v_par, v_perp], values)``: the extended grid (``v_perp`` a copy + of the ``v_par`` axis) and the composed distribution function, with a + trailing singleton component axis. + """ + x, vpar = f_grid[0], f_grid[1] + vperp = np.copy(vpar) + + x_cc = (x[:-1] + x[1:]) / 2 + vpar_cc = (vpar[:-1] + vpar[1:]) / 2 + vperp_cc = (vpar[:-1] + vpar[1:]) / 2 + + _, _, vperp_3D = np.meshgrid(x_cc, vpar_cc, vperp_cc, indexing="ij") + + F0 = f_values[..., 0] + G = f_values[..., 1] + T_m = t_over_m_values[..., 0] + + F1 = F0 - (G.transpose() / T_m).transpose() + + # Adding the np.newaxis allows the subsequent np.multiply (called when + # doing * on numpy arrays) to work. The arrays need to have the same + # number of axes, e.g. one cannot multiply (3, 3) and (3,) arrays but can + # multiply (3, 3) with (3, 1) or (1, 3). + F0, F1 = F0[..., np.newaxis], F1[..., np.newaxis] + # T_m gains two new axes here (F0/F1 gain only one above), one deeper than + # needed to broadcast against vperp_3D -- an extra, constant-along-itself + # trailing axis leaks into the returned array's shape. Preserved verbatim + # from src_bak/postgkyl/tools/laguerre_compose.py; pinned by + # tests/test_diagnostics_pkpm.py. + T_m = T_m[..., np.newaxis, np.newaxis] + + # Hardcoded for l=0, n=0,1 in + # https://drive.google.com/file/d/1548tLF9o7vyW3bkrsq6FvAMV-8XJvKtY/view + f = (F0 + F1 * (1 - vperp_3D**2 / 2 / T_m)) / (2 * np.pi * T_m) * np.exp( + -(vperp_3D**2) / 2 / T_m) + + f = f[..., np.newaxis] # Adding the component index + + return [x, vpar, vperp], f + + +# ---------------------------------------------------------------- GData verb +def laguerre_compose(distribution: "GDataState", + variables: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Compose PKPM Laguerre coefficients into a full distribution function. + + Reconstructs the full distribution function ``f(x, v_par, v_perp)`` from + the PKPM Laguerre expansion coefficients ``F0`` and ``G`` (stored as the + two components of ``distribution``) together with the PKPM + temperature-over-mass field carried in ``variables``. + + Args: + distribution: The two-component PKPM Laguerre expansion coefficients + ``F0(x, v_par)`` and ``G(x, v_par)``; must be NumPy-backed. + variables: The PKPM variables dataset providing T/m(x) (used as the + first component); must be NumPy-backed. + inplace: mutate and return ``distribution`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset holding the composed ``f(x, v_par, v_perp)``. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(distribution, "laguerre_compose", _REASON) + _require_field_domain(variables, "laguerre_compose", _REASON) + grid, values = _laguerre_compose(distribution.grid, distribution.values, + variables.values) + return distribution._result(grid, + values, + inplace=inplace, + tag=tag, + label=label) + + +# ------------------------------------------------------------------- loader +def load_pkpm(name: str, + species: str, + idx: "str | int", + poly_order: int, + *, + tag: str | None = None, + label: str | None = None) -> "GData": + """Load, interpolate, and frame-transform Gkeyll PKPM data. + + Loads the PKPM distribution (its two Laguerre coefficients ``F0``/``G``) + and its companion ``pkpm_vars`` file (whose component 3 is ``T/m`` and + components 0:3 are the bulk velocity ``(ux, uy, uz)``), interpolates both, + composes the full distribution function (:func:`laguerre_compose`), and + shifts it into the bulk-flow frame + (:func:`~postgkyl.diagnostics.vm.kinetic.transform_frame`). + + Args: + name: Root name (file prefix) of the simulation. + species: Species name. + idx: Frame/file number. + poly_order: Polynomial order of the DG representation. + tag: Optional tag for the resulting dataset. + label: Optional label for the resulting dataset. + + Returns: + A populated, interpolated, frame-transformed + :class:`~postgkyl.gdata.gdata.GData`. + """ + gf = GData(f"{name!s}-{species!s}_{idx!s}.gkyl", + basis_type="hybrid", + poly_order=poly_order) + gvars = GData(f"{name!s}-{species!s}_pkpm_vars_{idx!s}.gkyl", + basis_type="serendipity", + poly_order=poly_order) + + c_dim = gf.num_dims - 1 + + gf_interpolated = gf.interpolate() + gvars_interpolated = gvars.interpolate() + + t_over_m = gvars_interpolated.select(comp=3) + bulk_u = gvars_interpolated.select(comp="0:3") + + composed = laguerre_compose(gf_interpolated, t_over_m) + out = transform_frame(composed, bulk_u, cdim=c_dim) + + if tag is not None: + out.set_tag(tag) + if label is not None: + out.set_label(label) + return out + + +__all__ = ["laguerre_compose", "load_pkpm"] diff --git a/src/postgkyl/diagnostics/vm/__init__.py b/src/postgkyl/diagnostics/vm/__init__.py new file mode 100644 index 00000000..72287930 --- /dev/null +++ b/src/postgkyl/diagnostics/vm/__init__.py @@ -0,0 +1,5 @@ +"""Diagnostics for Gkeyll Vlasov and particle data.""" + +from . import kinetic, trajectory + +__all__ = ["kinetic", "trajectory"] diff --git a/src/postgkyl/diagnostics/vm/kinetic.py b/src/postgkyl/diagnostics/vm/kinetic.py new file mode 100644 index 00000000..4c3fb481 --- /dev/null +++ b/src/postgkyl/diagnostics/vm/kinetic.py @@ -0,0 +1,132 @@ +"""Distribution-function frame transform -- shift a particle distribution +function's velocity grid by a bulk velocity.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from ...gdatastate.guards import require_field_domain as _require_field_domain + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_REASON = "shifting the grid of raw DG coefficients has no basis-space meaning" + + +# --------------------------------------------------------- array-level math +def _transform_frame( + f_grid: list[np.ndarray], + f_values: np.ndarray, + u_values: np.ndarray, + c_dim: int, +) -> tuple[list[np.ndarray], np.ndarray]: + """Shift a distribution function to a different frame of reference. + + Shifts the velocity-space grid of a distribution function by a supplied + bulk velocity (a magnetic-field-direction shift is not yet supported). + + Args: + f_grid: Nodal coordinate arrays, one per configuration- and + velocity-space dimension (configuration dimensions first). + f_values: Particle distribution function values (unchanged by the + shift; only the velocity grid moves). + u_values: Bulk velocity array, ``num_dims - c_dim`` components, on the + configuration-space grid. + c_dim: Number of configuration-space dimensions. + + Returns: + ``(grid, values)``: a per-cell-shifted velocity grid (one nodal array + per dimension, matching the input's dimensionality) and the unchanged + distribution-function values. + """ + v_dim = len(f_grid) - c_dim + out_grid = np.meshgrid(*f_grid, indexing="ij") + + if c_dim == 1: + for v_idx in range(v_dim): + nx = f_grid[0].shape[0] + + ext_u = np.zeros(nx) + ext_u[:-1] += u_values[..., v_idx] + ext_u[1:] += u_values[..., v_idx] + ext_u[1:-1] = ext_u[1:-1] / 2 + + for i in range(nx): + out_grid[c_dim + v_idx][i, ...] += ext_u[i] + + elif c_dim == 2: + for v_idx in range(v_dim): + nx = f_grid[0].shape[0] + ny = f_grid[1].shape[0] + + ext_u = np.zeros((nx, ny)) + ext_u[:-1, :-1] += u_values[..., v_idx] + ext_u[1:, 1:] += u_values[..., v_idx] + ext_u[1:-1, 1:-1] = ext_u[1:-1, 1:-1] / 2 + + for i in range(nx): + for j in range(ny): + out_grid[c_dim + v_idx][i, j, ...] += ext_u[i, j] + + else: + for v_idx in range(v_dim): + nx = f_grid[0].shape[0] + ny = f_grid[1].shape[0] + nz = f_grid[2].shape[0] + + ext_u = np.zeros((nx, ny, nz)) + ext_u[:-1, :-1, :-1] += u_values[..., v_idx] + ext_u[1:, 1:, 1:] += u_values[..., v_idx] + ext_u[1:-1, 1:-1, 1:-1] = ext_u[1:-1, 1:-1, 1:-1] / 2 + + for i in range(nx): + for j in range(ny): + for k in range(nz): + out_grid[c_dim + v_idx][i, j, k, ...] += ext_u[i, j, k] + + return out_grid, f_values + + +# ---------------------------------------------------------------- GData verb +def transform_frame(distribution: "GDataState", + bulk: "GDataState", + *, + cdim: int, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Shift a distribution function to a moving frame of reference. + + Shifts the velocity-space grid of ``distribution`` by the local ``bulk`` + velocity so the distribution is expressed in the frame co-moving with + the bulk flow. The values are unchanged; only the velocity coordinates + are offset. Supports 1, 2, or 3 configuration-space dimensions. + + Args: + distribution: The particle distribution function to shift; must be + NumPy-backed. + bulk: The bulk (drift) velocity field; one component per velocity + dimension. Must be NumPy-backed. + cdim: Number of configuration-space dimensions. The remaining grid + axes are treated as velocity-space dimensions. + inplace: mutate and return ``distribution`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset with the same values on a velocity-shifted grid. + + Raises: + ValueError: if either input is native modal (gkyl-backed). + """ + _require_field_domain(distribution, "transform_frame", _REASON) + _require_field_domain(bulk, "transform_frame", _REASON) + grid, values = _transform_frame(distribution.grid, distribution.values, + bulk.values, cdim) + return distribution._result(grid, + values, + inplace=inplace, + tag=tag, + label=label) diff --git a/src/postgkyl/diagnostics/vm/trajectory.py b/src/postgkyl/diagnostics/vm/trajectory.py new file mode 100644 index 00000000..f166764d --- /dev/null +++ b/src/postgkyl/diagnostics/vm/trajectory.py @@ -0,0 +1,156 @@ +"""Particle-trajectory animation. + +Ported from ``src_bak/postgkyl/apps/trajectory.py``. Animates one or more +position (+ optional velocity) time series in 3-D. Typer options become +explicit keyword-only parameters; the old CLI's tag-indexed dataset stack +(``ctx.obj.data``) is replaced by passing the datasets directly. Saving is +the caller's choice: this returns the ``FuncAnimation`` object -- call +``.save(path)`` on it, or ``plt.show()`` after creating it to display it +live. + +Each dataset's ``grid[0]`` is expected to hold one time stamp per position +sample -- a Gkeyll dynvector's grid convention (``io/gkyl_reader.py``'s +``_read_t2_v1``: ``grid = [time]`` with ``len(time) == values.shape[0]``, +unlike a field file's ``num_cells + 1`` edges), the same convention +``src_bak`` read via ``dat.get_grid()[0]``. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.animation import FuncAnimation + +if TYPE_CHECKING: + from ...gdatastate.gdatastate import GDataState + +_COLORS = ("C0", "C1", "C2", "C3", "C4", "C5", "C6", "C7", "C8", "C9") + + +def _masked(coord: np.ndarray, lo: float | None, + hi: float | None) -> np.ndarray: + """Replace out-of-``[lo, hi]`` entries of ``coord`` with NaN, so they are + simply not drawn (masking, not clipping -- matches ``src_bak``).""" + out = coord + if lo is not None: + out = np.where(out > lo, out, np.nan) + if hi is not None: + out = np.where(out < hi, out, np.nan) + return out + + +def _update(i, ax, datasets, leap, no_velocity, xmin, xmax, ymin, ymax, zmin, + zmax): + """``FuncAnimation`` frame callback: redraw every dataset's trajectory up + to (and current position at) frame ``i``.""" + ax.cla() + t_idx = int(i * leap) + time = None + + for s, dataset in enumerate(datasets): + time = dataset.grid[0] + coords = dataset.values + color = _COLORS[s % len(_COLORS)] + + x = _masked(coords[:, 0], xmin, xmax) + y = _masked(coords[:, 1], ymin, ymax) + z = _masked(coords[:, 2], zmin, zmax) + + ax.plot(x, y, z, color=color) + ax.scatter(x[t_idx], y[t_idx], z[t_idx], color=color) + + if not no_velocity and dataset.num_comps == 6: + if t_idx + leap >= len(time): + dt = time[-1] - time[t_idx] + else: + dt = time[int(t_idx + leap)] - time[t_idx] + dx = coords[t_idx, 3] * dt + dy = coords[t_idx, 4] * dt + dz = coords[t_idx, 5] * dt + ax.plot([x[t_idx], x[t_idx] + dx], [y[t_idx], y[t_idx] + dy], + [z[t_idx], z[t_idx] + dz], + color=color) + + if time is not None: + ax.set_title(f"T: {time[t_idx]:.4e}") + ax.set_xlabel("$z_0$") + ax.set_ylabel("$z_1$") + ax.set_zlabel("$z_2$") + ax.set_xlim3d(xmin, xmax) + ax.set_ylim3d(ymin, ymax) + ax.set_zlim3d(zmin, zmax) + + +def trajectory( + *datasets: "GDataState", + fixaspect: bool = False, + interval: int = 100, + no_velocity: bool = False, + numframes: int | None = None, + xmin: float | None = None, + xmax: float | None = None, + ymin: float | None = None, + ymax: float | None = None, + zmin: float | None = None, + zmax: float | None = None, + elevation: float | None = None, + azimuth: float | None = None, +) -> FuncAnimation: + """Animate one or more particle trajectories in 3-D. + + Args: + datasets: One or more datasets, each holding a position (3-component, + ``x, y, z``) or position+velocity (6-component, + ``x, y, z, vx, vy, vz``) time series, with ``grid[0]`` one time stamp + per sample (the dynvector convention). + fixaspect: Enforce the same scaling on all three axes. + interval: Animation frame interval, in milliseconds. + no_velocity: Do not draw a velocity vector at the current position. + numframes: Number of animation frames; ``None`` uses one frame per + sample. When given, samples are subsampled evenly (by + ``floor(num_samples / numframes)``). + xmin: Optional lower x bound; outside points are masked. + xmax: Optional upper x bound; outside points are masked. + ymin: Optional lower y bound; outside points are masked. + ymax: Optional upper y bound; outside points are masked. + zmin: Optional lower z bound; outside points are masked. + zmax: Optional upper z bound; outside points are masked. + elevation: Initial 3-D elevation angle in degrees. + azimuth: Initial 3-D azimuth angle in degrees. + + Returns: + The ``FuncAnimation``. + + Raises: + ValueError: if no datasets are given. + """ + if not datasets: + raise ValueError("trajectory() requires at least one dataset.") + + fig = plt.figure() + ax = fig.add_subplot(111, projection="3d") + + num_pos = int(datasets[0].num_cells[0]) + leap = 1 + if numframes: + leap = int(math.floor(num_pos / numframes)) + num_pos = int(numframes) + + anim = FuncAnimation(fig, + _update, + num_pos, + fargs=(ax, datasets, leap, no_velocity, xmin, xmax, ymin, + ymax, zmin, zmax), + interval=interval) + + ax.view_init(elev=elevation, azim=azimuth) + if fixaspect: + # Equal-scale 3-D axes: modern Matplotlib's Axes3D takes a box aspect + # ratio (`set_box_aspect`), not the numeric `aspect=` src_bak passed to + # `plt.setp` (that spelling only ever worked for 2-D axes). + ax.set_box_aspect((1.0, 1.0, 1.0)) + + return anim diff --git a/src/postgkyl/gdata/__init__.py b/src/postgkyl/gdata/__init__.py new file mode 100644 index 00000000..60198eea --- /dev/null +++ b/src/postgkyl/gdata/__init__.py @@ -0,0 +1,13 @@ +"""The fluent API surface: the public ``GData``, ``load``, ``GDataGroup``, +and the module-level multi-dataset verbs (``collect``/``evaluate``/``relchange``/ +``animate``/``plotly_animate``/``sort``).""" + +from .gdata import GData +from .load import load +from .gdatagroup import GDataGroup +from .verbs import animate, collect, evaluate, plot, plotly_animate, relchange, sort + +__all__ = [ + "GData", "load", "GDataGroup", "collect", "evaluate", "relchange", "plot", + "animate", "plotly_animate", "sort" +] diff --git a/src/postgkyl/gdata/gdata.py b/src/postgkyl/gdata/gdata.py new file mode 100644 index 00000000..d5b1ac6c --- /dev/null +++ b/src/postgkyl/gdata/gdata.py @@ -0,0 +1,241 @@ +"""``GData`` -- the fluent surface (the FLUENT API layer). + +A thin subclass of the verb-less :class:`~postgkyl.gdatastate.gdatastate.GDataState` +container that adds the fluent verb methods and the computing operators. Because +this module sits *above* ``operations``/``render``/``io``, it imports them with plain +top-level imports -- there is **no import cycle and no lazy import anywhere**. + +Inherited from the container (pure state readers): ``info``, ``__array__``, +``__repr__``/``__str__``, all shape properties, ``copy``/``_result``. +""" + +from __future__ import annotations + +from glob import has_magic +import operator + +import numpy as np + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl import operations, io +from postgkyl.cli_spec import hidden + +from .gdatagroup import GDataGroup + + +class GData(GDataState): + """Fluent dataset: ``pg.load(...).interpolate().select(z0=0.0).plot()``.""" + + # ------------------------------------------------------- data lifecycle + def load(self, + file_name: str, + *, + tag: str | None = None, + label: str | None = None, + ctx: dict | None = None, + value_form: str | None = None, + basis_type: str | None = None, + poly_order: int | None = None, + **read_kwargs) -> "GData": + """Load one file into this dataset in place and return ``self``. + + This is the two-step counterpart of constructing ``GData(file_name)``:: + + data = GData() + data.load(file_name).local_poly().plot() + + The read is atomic with respect to this object: if it fails, the current + grid, values, context, and filename are left unchanged. A pristine empty + dataset's existing ``ctx`` seeds the read unless ``ctx`` is supplied + explicitly; reloading an already populated dataset starts from a fresh + context so metadata from the previous file cannot leak into the new one. + The dataset's existing tag and custom label are preserved unless ``tag`` + or ``label`` is passed. + + This method accepts one literal filename. Use :func:`postgkyl.load` for + shell-style glob patterns, which produce a ``GDataGroup`` rather than one + dataset. + """ + file_name = str(file_name) + if not file_name: + raise ValueError("GData.load() requires a non-empty filename.") + if has_magic(file_name): + raise ValueError( + "GData.load() accepts one literal filename; use pg.load(pattern) " + "to load a glob as a GDataGroup.") + + if ctx is None and self._grid is None and self._values is None: + load_ctx = self.ctx + else: + load_ctx = ctx + + # Construct through GDataState so this follows exactly the same reader and + # metadata-defaulting path as GData(file_name). Nothing on ``self`` is + # changed until construction succeeds. + loaded = GDataState(file_name, + ctx=load_ctx, + value_form=value_form, + basis_type=basis_type, + poly_order=poly_order, + **read_kwargs) + self._grid = loaded._grid + self._values = loaded._values + self.ctx = loaded.ctx + self._file_name = loaded._file_name + self._label = loaded._label + if tag is not None: + self._tag = tag + if label is not None: + self._custom_label = label + return self + + # ---------------------------------------------------------- fluent verbs + # These are ordinary class-body aliases, not wrappers or runtime setattr + # calls. Python binds the leading dataset argument as ``self``. The alias + # keeps the signature, annotations, docstring, command metadata, and + # implementation in one canonical function while remaining discoverable by + # static language servers such as VS Code/Pylance. + interpolate = operations.interpolate + local_poly = operations.local_poly + gk_rz = operations.gyrokinetics.gk_rz + gk_fluxsurf = operations.gyrokinetics.gk_fluxsurf + select = operations.select + integrate = operations.integrate + average = operations.average + eval_at_coord_proj = operations.eval_at_coord_proj + fft = operations.fft + magsq = operations.magsq + mask = operations.mask + extract_input = operations.extract_input + fit = operations.fit + growth = operations.growth + differentiate = operations.differentiate + map = operations.map + apply = operations.apply + save = io.save + plot = operations.plot + plotly = operations.plotly + pyvista = operations.pyvista + + # ``info`` is inherited from GDataState (a pure state reader). + + # ----------------------------------------------------------- modal verbs + # Explicit spellings of the weak algebra (the * and / operators dispatch to + # the same Gkeyll kernels when both operands are modal). + def mul(self, other) -> "GData": + """Weak (DG) multiply -- runs inside Gkeyll on modal data.""" + return operations.arithmetic.binary(operator.mul, self, other) + + def div(self, other) -> "GData": + """Weak (DG) divide -- runs inside Gkeyll on modal data.""" + return operations.arithmetic.binary(operator.truediv, self, other) + + # --------------------------------------------- value_form changes (explicit) + # Conversions never happen implicitly -- these verbs are the only doorway + # between the modal / nodal / quadrature value_forms (all gkyl-native). + def to_modal(self, **kwargs) -> "GData": + """Convert to modal coefficients (exact from nodal; projection from quad).""" + return operations.represent(self, to="modal", **kwargs) + + def to_nodal(self, **kwargs) -> "GData": + """Convert to values at the basis nodes (exact, invertible).""" + return operations.represent(self, to="nodal", **kwargs) + + def to_quad(self, num_quad: int | None = None, **kwargs) -> "GData": + """Convert to values at Gauss–Legendre points (default ``p+1`` per dim).""" + return operations.represent(self, to="quad", num_quad=num_quad, **kwargs) + + # ------------------------------------------------- field-domain analysis + def val2coord(self, + *, + x: str, + y: str, + periodic: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataGroup": + """Build new (x, y) datasets from DynVector columns (see ``operations.val2coord``). + + Wraps the ``operations`` verb's (verb-less) ``core.GDataStateGroup`` result in a + fluent :class:`~postgkyl.gdata.gdatagroup.GDataGroup` so the chain keeps going, + e.g. ``d.val2coord(x='0', y='1:3')[0].plot()``. + """ + return GDataGroup( + operations.val2coord(self, + x=x, + y=y, + periodic=periodic, + tag=tag, + label=label)) + + # Note: no fluent ``grid`` method. ``GData.grid`` (inherited from + # GDataState) is the axis-edge-array property that most of ``operations`` reads + # via plain attribute access (``data.grid``); a same-named verb method + # would shadow it for every GData instance and silently break every other + # verb. ``operations.grid`` (the "turn a dataset's grid into a dataset of + # coordinates" verb) is reachable as ``postgkyl.operations.grid(data, ...)`` -- + # src_bak's GData carried the identical exception with the identical + # reasoning (src_bak/postgkyl/data/gdata.py:1258-1259). + + # ------------------------------------------------------ binary operators + def __add__(self, o): + return operations.arithmetic.binary(operator.add, self, o) + + def __sub__(self, o): + return operations.arithmetic.binary(operator.sub, self, o) + + def __mul__(self, o): + return operations.arithmetic.binary(operator.mul, self, o) + + def __truediv__(self, o): + return operations.arithmetic.binary(operator.truediv, self, o) + + def __pow__(self, o): + return operations.arithmetic.binary(operator.pow, self, o) + + def __radd__(self, o): + return operations.arithmetic.binary(operator.add, o, self) + + def __rsub__(self, o): + return operations.arithmetic.binary(operator.sub, o, self) + + def __rmul__(self, o): + return operations.arithmetic.binary(operator.mul, o, self) + + def __rtruediv__(self, o): + return operations.arithmetic.binary(operator.truediv, o, self) + + def __rpow__(self, o): + return operations.arithmetic.binary(operator.pow, o, self) + + # ----------------------------------------------------------------- unary + def __neg__(self): + return operations.arithmetic.binary(operator.mul, self, -1.0) + + def __abs__(self): + return operations.arithmetic.apply_ufunc(np.absolute, "__call__", self) + + def __pos__(self): + return self.clone() + + # --------------------------------------------------------- NumPy interop + __array_priority__ = 100 # ndarray defers to us in mixed ndarray·GData ops + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + """Apply NumPy ufuncs while preserving pointwise dataset metadata. + + Pointwise calls such as ``np.sqrt``/``np.add`` return a GData carrying the + grid/ctx; reductions such as ``np.max``/``np.sum`` return NumPy results. + """ + return operations.arithmetic.apply_ufunc(ufunc, method, *inputs, **kwargs) + + +for _name, _reason in { + "load": "the canonical loader is postgkyl.load", + "mul": "Python operators are not stringly exposed as commands", + "div": "Python operators are not stringly exposed as commands", + "to_modal": "representation shortcuts remain Python-only", + "to_nodal": "representation shortcuts remain Python-only", + "to_quad": "representation shortcuts remain Python-only", + "val2coord": "the functional operation owns this exceptional group result", +}.items(): + hidden(_reason)(GData.__dict__[_name]) diff --git a/src/postgkyl/gdata/gdatagroup.py b/src/postgkyl/gdata/gdatagroup.py new file mode 100644 index 00000000..16cce782 --- /dev/null +++ b/src/postgkyl/gdata/gdatagroup.py @@ -0,0 +1,189 @@ +"""``GDataGroup`` -- the fluent group container over +``gdatastate.GDataStateGroup``. + +Mirrors how :class:`~postgkyl.gdata.gdata.GData` adds the fluent verb methods +on top of the verb-less :class:`~postgkyl.gdatastate.gdatastate.GDataState`: this class +adds *broadcasting* verbs on top of the verb-less +:class:`~postgkyl.gdatastate.gdatastategroup.GDataStateGroup`, without duplicating a single +verb body. + +Contract +-------- +Any attribute name that is not defined on this class itself (and does not +start with ``_``) is resolved by :meth:`__getattr__`, by looking it up on +every member, in order (**broadcasting**): + +- If the attribute is a *verb method* on every member, calling the broadcast + invokes that method on each member with the same arguments. If every + member's result is a ``GDataState`` (or subclass), the results are wrapped + in a *new* group of the caller's own concrete class, so chains stay fluent: + ``group.interpolate().select(z0=0.0)``. Otherwise -- a terminal verb whose result is + not a dataset (``.write()`` -> one path per member, full ``.integrate()`` -> + one value per member, ``.extract_input()`` -> one string per member, ...) + -- a plain ``list`` of the per-member results is returned, in member order. +- If the attribute is a *non-callable* value on every member (a property such + as ``num_dims`` or ``backend``), it resolves immediately to a plain + ``list`` of the per-member values, in member order -- no closure, no + call needed. +- Attribute names starting with ``_`` are never broadcast (raises + ``AttributeError``), so private/dunder probes and pickling machinery are + unaffected. An attribute missing from any member also raises + ``AttributeError`` immediately, at access time. + +The explicit methods below either combine/reorder the members or have +collection-wide terminal semantics. In particular, ``plot`` draws every +member on one figure, so a multiblock family remains one physical field. +``plotly`` is still an ordinary per-dataset verb and therefore broadcasts. + +``operations.grid`` has no fluent spelling anywhere (not on ``GData``, so not +broadcast here either) -- see ``api/gdata.py`` for why. + +``load`` is also explicit rather than broadcast: it is the group's lifecycle +method, appending newly loaded member(s) to this group and returning ``self`` +so an initially empty group can be assembled fluently. +""" + +from __future__ import annotations + +from postgkyl import operations +from postgkyl.gdatastate.gdatastategroup import GDataStateGroup +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.cli_spec import hidden + +from . import verbs + + +class GDataGroup(GDataStateGroup): + """A group whose members' fluent verbs broadcast over the whole group.""" + + # ------------------------------------------------------- data lifecycle + def load(self, + file_name: str, + *, + tag: str = "default", + label: str = "", + ctx: dict | None = None, + value_form: str | None = None, + basis_type: str | None = None, + poly_order: int | None = None, + **read_kwargs) -> "GDataGroup": + """Load and append one file (or a glob) and return this group. + + Successive calls accumulate members, enabling chains such as:: + + group = GDataGroup() + group.load(frame0).load(frame1).local_poly().collect().plot() + + Loading completes before this group is changed, so a failed read leaves + its existing members untouched. A glob appends every match in natural + filename order, with the same options and errors as :func:`postgkyl.load`. + """ + # Same-layer import at call time avoids the construction-time cycle: + # gdata.load builds GDataGroup results, while GData imports this class for + # fluent group-returning verbs. + from .load import load as load_data + + loaded = load_data(file_name, + tag=tag, + label=label, + ctx=ctx, + value_form=value_form, + basis_type=basis_type, + poly_order=poly_order, + **read_kwargs) + if isinstance(loaded, GDataStateGroup): + additions = loaded.datasets + else: + additions = [loaded] + self._datasets.extend(additions) + return self + + def __getattr__(self, name: str): + if name.startswith("_"): + raise AttributeError(name) + values = [getattr(member, name) for member in self._datasets] + if values and not all(callable(v) for v in values): + return values + + def broadcast(*args, **kwargs): + results = [v(*args, **kwargs) for v in values] + if results and all(isinstance(r, GDataState) for r in results): + return type(self)(results) + return results + + return broadcast + + # ------------------------------------------------------- combining (typed) + # Overridden (not inherited) so the result stays the caller's concrete + # subclass, mirroring GDataState._result's ``type(self)`` trick. + def with_(self, *others) -> "GDataGroup": + """Return a new group (same concrete class) with ``others`` appended.""" + return type(self)(self._datasets + list(others)) + + __and__ = with_ + + def sort(self, *, reverse: bool = False) -> "GDataGroup": + """Return a naturally filename-sorted group (see ``operations.sort``).""" + return type(self)(verbs.sort(*self._datasets, reverse=reverse)) + + def __getitem__(self, index): + """Index or slice; a slice returns a group of the same concrete class.""" + result = self._datasets[index] + return type(self)(result) if isinstance(index, slice) else result + + # ------------------------------------------------------- terminal (typed) + def info(self, *, no_header: bool = False) -> list: + """Summarize every member (see ``operations.info``); returns a list of strings.""" + return operations.info(*self._datasets, no_header=no_header) + + # Binding the canonical variadic function passes this iterable group as its + # first input; render.plot flattens it and keeps group calls on one figure. + plot = verbs.plot + + def collect(self, + *, + sumdata: bool = False, + period: float | None = None, + offset: float = 0.0, + chunk: int | None = None, + tag: str | None = None, + label: str | None = None): + """Combine the members into one dataset along a time axis (see + ``api.verbs.collect``). Returns a list of datasets when ``chunk`` is given.""" + return verbs.collect(*self._datasets, + sumdata=sumdata, + period=period, + offset=offset, + chunk=chunk, + tag=tag, + label=label) + + def evaluate(self, + chain: str, + *, + tag: str | None = None, + label: str | None = None): + """Evaluate an RPN expression over the members (see ``api.verbs.evaluate``).""" + return verbs.evaluate(chain, *self._datasets, tag=tag, label=label) + + def animate(self, **kwargs): + """Animate the members, one frame each (see ``api.verbs.animate``).""" + return verbs.animate(self._datasets, **kwargs) + + def plotly_animate(self, **kwargs): + """Animate the members with Plotly, one frame each (see ``api.verbs.plotly_animate``).""" + return verbs.plotly_animate(self._datasets, **kwargs) + + +for _name in ( + "load", + "with_", + "sort", + "info", + "collect", + "evaluate", + "animate", + "plotly_animate", +): + hidden("the functional callable is the canonical command source")( + GDataGroup.__dict__[_name]) diff --git a/src/postgkyl/gdata/load.py b/src/postgkyl/gdata/load.py new file mode 100644 index 00000000..521a643d --- /dev/null +++ b/src/postgkyl/gdata/load.py @@ -0,0 +1,124 @@ +"""``pg.load`` -- load one file or a glob of files into the fluent API.""" + +from __future__ import annotations + +from glob import glob, has_magic +from typing import Annotated, Literal + +from postgkyl import operations +from postgkyl.cli_spec import ( + CommandSpec, + Execution, + KeyValue, + Section, + command, +) +from postgkyl.gdata.gdata import GData +from postgkyl.gdata.gdatagroup import GDataGroup + + +@command(CommandSpec(Section.UTILITY, Execution.LOAD)) +def load( + file_name: str, + *, + tag: str = "default", + label: str = "", + ctx: Annotated[dict[str, str] | None, KeyValue()] = None, + value_form: Literal["modal", "nodal", "quad"] | None = None, + basis_type: str | None = None, + poly_order: int | None = None, + z0: str | None = None, + z1: str | None = None, + z2: str | None = None, + z3: str | None = None, + z4: str | None = None, + z5: str | None = None, + component: str | None = None, + read_options: Annotated[dict[str, str] | None, + KeyValue()] = None, +) -> GData | GDataGroup: + """Read Gkeyll output into a fluent ``GData`` or ``GDataGroup``. + + ``pg.load('elc_M0_0.gkyl').interpolate().select(z0=0.0).plot()`` + + Shell-style glob patterns (``*``, ``?``, and ``[]``) load every matching + file into a :class:`GDataGroup`, naturally ordered by filename. The group + broadcasts per-dataset verbs and supplies the multi-dataset ``collect`` + verb, so a frame series can be loaded and stacked in one chain:: + + pg.load('elc_M0_*.gkyl').interpolate().collect().plot() + + A pattern always returns a group, even if it matches only one file. A + literal filename retains the original single-``GData`` return type. + + ``basis_type``, ``poly_order``, and ``value_form`` are properties of the + data itself, fixed here at load time (from the file's header metadata, or + the override below) -- no downstream verb (``interpolate``, ``average``, + ...) ever re-specifies them; they always read ``ctx["basis_type"]``/ + ``ctx["poly_order"]``/``ctx["value_form"]`` off the loaded dataset. + + ``value_form`` overrides the ``"modal"``/``"nodal"``/``"quad"`` tag the + file's header metadata would otherwise imply -- for files whose writer + stamps DG basis metadata even though the stored values are already point + values (e.g. a per-cell diagnostic like a CFL rate), not modal coefficients. + + ``basis_type`` overrides the ``"basis_type"`` (e.g. ``"serendipity"``, + ``"tensor"``, ``"gkhybrid"``) the file's header metadata would otherwise + imply -- for files with no basis metadata at all, or metadata that + mislabels the basis actually used. Setting it also defaults ``value_form`` + to ``"modal"`` (unless ``value_form`` is given too), so downstream verbs + that read ``ctx["basis_type"]`` resolve the right basis. + + ``poly_order`` overrides the ``"poly_order"`` the file's header metadata + would otherwise imply. It is independent of ``basis_type``/``value_form`` -- + passing it alone corrects only the polynomial order and asserts nothing + about whether the dataset is modal. + + Args: + file_name: Literal filename or shell-style glob pattern to load. + tag: Tag assigned to every loaded dataset. + label: Optional display label assigned to every loaded dataset. + ctx: Initial metadata as repeated key/value entries. + value_form: Stored representation of the loaded values. + basis_type: DG basis name overriding file metadata. + poly_order: Polynomial order overriding file metadata. + z0: Partial-load selector for coordinate direction 0. + z1: Partial-load selector for coordinate direction 1. + z2: Partial-load selector for coordinate direction 2. + z3: Partial-load selector for coordinate direction 3. + z4: Partial-load selector for coordinate direction 4. + z5: Partial-load selector for coordinate direction 5. + component: Partial-load component selector. + read_options: Reader-specific options as repeated key/value entries. + """ + file_name = str(file_name) + read_kwargs = dict(read_options or {}) + axes = (z0, z1, z2, z3, z4, z5) + if any(value is not None for value in axes): + read_kwargs["axes"] = axes + if component is not None: + read_kwargs["comp"] = component + if has_magic(file_name): + matches = glob(file_name) + if not matches: + raise FileNotFoundError(f"No files match pattern: '{file_name}'") + datasets = [ + GData(match, + tag=tag, + label=label, + ctx=ctx, + value_form=value_form, + basis_type=basis_type, + poly_order=poly_order, + **read_kwargs) for match in matches + ] + return GDataGroup(operations.sort(datasets)) + + return GData(file_name, + tag=tag, + label=label, + ctx=ctx, + value_form=value_form, + basis_type=basis_type, + poly_order=poly_order, + **read_kwargs) diff --git a/src/postgkyl/gdata/verbs.py b/src/postgkyl/gdata/verbs.py new file mode 100644 index 00000000..d0c6c1cc --- /dev/null +++ b/src/postgkyl/gdata/verbs.py @@ -0,0 +1,30 @@ +"""Module-level fluent verbs -- the multi-dataset verbs that have no single +``self``. + +``collect``, ``evaluate``, ``relchange``, ``plot``, ``animate``, and +``plotly_animate`` each combine *several* datasets into one result (or, for +``plot``/``animate``/``plotly_animate``, into one figure/animation); ``sort`` +reorders several datasets rather than combining them. None of these can be one dataset's method the +way ``interpolate``/``select``/``fft``/... are on +:class:`~postgkyl.gdata.gdata.GData`. Each is a direct alias to the canonical +callable exposed through :mod:`postgkyl.operations`, so the functional spelling +(``postgkyl.collect(a, b)``) and this module-level fluent spelling can never +drift apart. :class:`~postgkyl.gdata.gdatagroup.GDataGroup` re-uses these same +functions for its own ``sort``/``collect``/``evaluate``/``plot``/``animate``/ +``plotly_animate`` methods. +""" + +from __future__ import annotations + +from postgkyl import operations + +# Exact aliases: implementation, signature, annotations, docstring, and command +# metadata remain on the canonical callable (``render`` owns ``plot``, +# ``animate``, and ``plotly_animate``). +collect = operations.collect +sort = operations.sort +evaluate = operations.evaluate +relchange = operations.relchange +animate = operations.animate +plotly_animate = operations.plotly_animate +plot = operations.plot diff --git a/src/postgkyl/gdatastate/__init__.py b/src/postgkyl/gdatastate/__init__.py new file mode 100644 index 00000000..66193174 --- /dev/null +++ b/src/postgkyl/gdatastate/__init__.py @@ -0,0 +1,15 @@ +"""The object-model layer: the verb-less ``GDataState`` container.""" + +from .gdatastate import GDataState +from .collection import flatten_datasets, group_blocks, group_frames +from .gdatastategroup import GDataStateGroup +from .materialize import materialize_point_values + +__all__ = [ + "GDataState", + "flatten_datasets", + "group_blocks", + "group_frames", + "GDataStateGroup", + "materialize_point_values", +] diff --git a/src/postgkyl/gdatastate/collection.py b/src/postgkyl/gdatastate/collection.py new file mode 100644 index 00000000..7d7625b7 --- /dev/null +++ b/src/postgkyl/gdatastate/collection.py @@ -0,0 +1,106 @@ +"""Helpers for collections of datasets (shared by the multi-dataset verbs). + +Lives in ``gdatastate`` because it is generic plumbing over the container type and is +needed by both ``render`` (``pg.plot(a, b)``) and ``operations`` (``pg.info(a, b)``) -- +both of which already depend on ``gdatastate``. Keeping it here avoids duplicating the +flatten in two layers or stranding it in the facade. +""" + +from __future__ import annotations + +from .gdatastate import GDataState + + +def flatten_datasets(items) -> list: + """Flatten nested lists/tuples/groups of datasets into a single flat list. + + Lets the multi-dataset entry points accept either ``f(a, b)`` or ``f([a, b])`` + (and nested combinations, including a ``GDataStateGroup`` wherever a dataset is + expected). Recursion is on any iterable, not just ``list``/``tuple`` -- this is + what lets a nested ``gdatastate.gdatastategroup.GDataStateGroup`` flatten correctly without this + module importing that one (it needs no type check, only that groups are + iterable). Strings pass through whole (never iterated character-by-character); + non-dataset, non-iterable items also pass through so the downstream consumer + can raise a clear, contextual error. + """ + out = [] + for it in items: + if isinstance(it, GDataState): + out.append(it) + elif isinstance(it, (str, bytes)): + out.append(it) + elif hasattr(it, "__iter__"): + out.extend(flatten_datasets(it)) + else: + out.append(it) + return out + + +def _family_key(data) -> tuple | None: + """The key identifying "the same field" across a multiblock decomposition, + or ``None`` for a dataset that is not part of one. + + Built from the identity ``GDataState`` stamps at load time (``sim``, + ``quantity``, ``frame`` -- see ``io.naming``) plus the dataset's ``tag``, + so two differently-tagged results of the same source file (e.g. the raw + load and a ``gk_rz`` projection of it) never merge. ``block`` is + deliberately absent: it is what family members differ by. + + Returning ``None`` for single-block data is the property that keeps every + pre-existing pipeline byte-identical -- with no ``_b`` in the file + names, every dataset is its own family. + """ + if not isinstance(data, GDataState) or data.ctx.get("block") is None: + return None + return (data.tag, data.ctx.get("sim"), data.ctx.get("quantity"), + data.ctx.get("frame")) + + +def group_blocks(datasets) -> list[list]: + """Partition datasets into **block families**: one field's blocks together. + + A family is the set of datasets that agree on ``(tag, sim, quantity, + frame)`` and differ only in ``ctx["block"]`` -- i.e. the pieces of one + field on a decomposed domain, which terminal verbs (``plot``, ``animate``) + should treat as a single thing to draw. Datasets with no block index are + each returned as their own singleton family, so single-block input maps + 1:1 onto the ungrouped list it was before. + + Args: + datasets: Datasets, groups, or nested iterables of them (flattened via + :func:`flatten_datasets`). + + Returns: + A list of lists, in first-appearance order; each family is sorted by + ascending block index. + """ + families: dict = {} + out: list[list] = [] + for data in flatten_datasets(datasets): + key = _family_key(data) + if key is None: + out.append([data]) + continue + if key not in families: + families[key] = [] + out.append(families[key]) + families[key].append(data) + for family in families.values(): + family.sort(key=lambda d: int(d.ctx["block"])) + return out + + +def group_frames(datasets) -> list[list]: + """Group datasets by authoritative ``ctx['frame']`` metadata. + + Known frames are returned in ascending order. Datasets without a frame + stay together in one trailing group. + """ + groups: dict[int | None, list] = {} + for data in flatten_datasets(datasets): + frame = data.ctx.get("frame") + groups.setdefault(int(frame) if frame is not None else None, + []).append(data) + known = sorted(frame for frame in groups if frame is not None) + return [groups[frame] + for frame in known] + ([groups[None]] if None in groups else []) diff --git a/src/postgkyl/gdatastate/gdatastate.py b/src/postgkyl/gdatastate/gdatastate.py new file mode 100644 index 00000000..f830d2b5 --- /dev/null +++ b/src/postgkyl/gdatastate/gdatastate.py @@ -0,0 +1,502 @@ +"""``GDataState`` -- the verb-less data container (the CONTAINER layer). + +Holds a Gkeyll dataset: a nodal ``grid`` (list of 1-D edge arrays) plus values +in one of **two backends** -- the two-domain lifecycle of REFACTOR_GKEYLL_FFI.md: + +- ``backend == "gkyl"``: modal DG coefficients held as a native + :class:`~postgkyl.gpython.array.GkylArray`. Gkeyll owns the memory and all math + on it (weak ops, coefficient lin-combs, integrate). ``values`` exposes a + read-only NumPy *view* for inspection; ``__array__`` refuses (interpolate first). +- ``backend == "numpy"``: post-``interpolate`` (or never-modal) values as a plain + ``np.ndarray`` -- the field domain, where all NumPy math applies. + +It constructs itself by delegating to the :mod:`postgkyl.io` leaf and exposes +only *state*. Crucially it imports **nothing upward** (no ``operations``/``render``/ +``api``). The fluent verb methods and the computing operators live on the +:class:`postgkyl.gdata.gdata.GData` subclass, one layer up. That is what keeps +the dependency graph a strict, cycle-free DAG -- see HIERARCHY_2.md / HIERARCHY_3.md. +""" + +from __future__ import annotations + +import numbers +import warnings +from typing import Tuple + +import numpy as np + +from postgkyl import io # leaf layer (below); top-level import -- never a cycle +from postgkyl import gpython # foreign floor (below): GkylArray backend type + + +class GDataState: + """Storage + metadata for one dataset. No verbs; no upward imports.""" + + def __init__(self, + file_name: str = "", + *, + ctx: dict | None = None, + tag: str = "default", + label: str = "", + value_form: str | None = None, + basis_type: str | None = None, + poly_order: int | None = None, + **read_kwargs): + self._grid: list | None = None + self._values: np.ndarray | gpython.GkylArray | None = None + self.ctx: dict = {} + if ctx: + self.ctx.update(ctx) + self._tag = tag + self._label = "" + self._custom_label = label + self._file_name = str(file_name) + self.color = None + + if self._file_name: + self._grid, self._values = io.read(self._file_name, + self.ctx, + value_form=value_form, + basis_type=basis_type, + poly_order=poly_order, + **read_kwargs) + self._stamp_output_name() + # A dynvector/diagnostic file (no "cells" in ctx: no reader ever stamps + # one without a spatial grid, e.g. a dynvector time series) has no DG + # basis to speak of -- basis_type/poly_order/value_form genuinely don't + # apply, so no defaulting is needed here. + if self.ctx.get("cells") is not None: + defaulted = [] + if self.ctx.get("basis_type") is None: + self.ctx["basis_type"] = "serendipity" + defaulted.append("basis_type") + if self.ctx.get("poly_order") is None: + self.ctx["poly_order"] = 0 + defaulted.append("poly_order") + if "value_form" not in self.ctx: + self.ctx["value_form"] = "nodal" + defaulted.append("value_form") + if defaulted: + warnings.warn( + f"{self._file_name}:\n" + f"{', '.join(defaulted)} not resolvable (not present in the " + "file header, and not given explicitly); defaulting to " + "basis_type='serendipity', poly_order=0, value_form='nodal' " + "(p0 -- one point per cell, at the cell center). Pass " + "basis_type=/poly_order=/value_form=... explicitly if this " + "is wrong.", + stacklevel=2) + + # -------------------------------------------------------------- identity + def _stamp_output_name(self) -> None: + """Record the file's Gkeyll *identity* (sim, block, quantity, frame) in + ``ctx``, parsed once from its path by :mod:`postgkyl.io.naming`. + + Header metadata wins: ``setdefault`` never overwrites a ``frame`` (or + anything else) a reader already read out of the file itself. Because + ``clone`` copies ``ctx``, the identity survives every verb, so a + multiblock family is still recognizable after ``interpolate``/``gk_rz`` + -- which is what lets terminal verbs draw one field's blocks together + (see ``gdatastate.collection.group_blocks``). + """ + name = io.parse_output_name(self._file_name) + if name is None: + return + self.ctx.setdefault("sim", name.sim) + self.ctx.setdefault("block", name.block) + self.ctx.setdefault("quantity", name.quantity) + if name.frame is not None: + self.ctx.setdefault("frame", name.frame) + + @property + def output_name(self): + """This dataset's parsed source-file identity (:class:`postgkyl.io.OutputName`), + or ``None`` when it was never read from disk.""" + return io.parse_output_name(self._file_name) + + # ------------------------------------------------------------------ tags + def get_tag(self) -> str: + """Return the short identifier used to select this dataset.""" + return self._tag + + def set_tag(self, tag: str = "") -> None: + """Replace the dataset tag when ``tag`` is nonempty.""" + if tag: + self._tag = tag + + tag = property(get_tag, set_tag) + + def get_label(self) -> str: + """Return the custom label, falling back to the generated label.""" + return self._custom_label or self._label + + def set_label(self, label: str) -> None: + """Set the generated display label.""" + self._label = label + + label = property(get_label, set_label) + + @property + def file_name(self) -> str: + """Source file path this dataset was loaded from ("" if it was never + read from disk, e.g. a verb's freshly-computed result).""" + return self._file_name + + # ------------------------------------------------------------- shape info + def get_num_cells(self) -> np.ndarray: + """Return the cell count in each spatial dimension.""" + if self.ctx.get("cells") is not None: + return np.asarray(self.ctx["cells"]) + if isinstance(self._values, np.ndarray): + return np.array(self._values.shape[:-1], dtype=np.int64) + return np.array([], dtype=np.int64) + + num_cells = property(get_num_cells) + + def get_num_comps(self) -> int: + """Return the number of physical components per cell.""" + if self.ctx.get("num_comps"): + return int(self.ctx["num_comps"]) + if isinstance(self._values, gpython.GkylArray): + return self._values.ncomp + if self._values is not None: + return int(self._values.shape[-1]) + return 0 + + num_comps = property(get_num_comps) + + def get_num_dims(self) -> int: + """Return the number of spatial dimensions.""" + if self.ctx.get("cells") is not None: + return len(self.ctx["cells"]) + if isinstance(self._values, np.ndarray): + return int(self._values.ndim - 1) + return 0 + + num_dims = property(get_num_dims) + + def get_bounds(self) -> Tuple[np.ndarray, np.ndarray]: + """Return arrays containing the lower and upper spatial bounds.""" + if "lower" in self.ctx and "upper" in self.ctx: + return np.asarray(self.ctx["lower"]), np.asarray(self.ctx["upper"]) + if self._grid is not None: + num_dims = self.get_num_dims() + lo = np.array([self._grid[d].min() for d in range(num_dims)]) + up = np.array([self._grid[d].max() for d in range(num_dims)]) + return lo, up + return None, None + + bounds = property(get_bounds) + + def get_grid_type(self) -> str: + """Return the grid classification, defaulting to ``"uniform"``.""" + return self.ctx.get("grid_type", "uniform") + + # --------------------------------------------------------- grid / values + def get_grid(self) -> list: + """Return the coordinate array for each spatial dimension.""" + return self._grid + + def set_grid(self, grid: list) -> None: + """Replace the coordinate arrays and update their stored bounds.""" + self._grid = grid + # ``len(grid)`` (not ``get_num_dims()``) on purpose: for a gkyl-backed + # dataset, num_dims reads ctx["cells"], which a dimension-reducing verb + # (e.g. ``average``) updates via ``_result``'s ctx_updates -- AFTER + # ``push`` calls this method. Deriving straight from the just-given grid + # avoids depending on that update having landed yet. + num_dims = len(grid) + self.ctx["lower"] = np.array([grid[d].min() for d in range(num_dims)]) + self.ctx["upper"] = np.array([grid[d].max() for d in range(num_dims)]) + + grid = property(get_grid, set_grid) + + @property + def backend(self) -> str: + """``"gkyl"`` (native modal storage) or ``"numpy"`` (field domain).""" + return "gkyl" if isinstance(self._values, gpython.GkylArray) else "numpy" + + @property + def native(self) -> gpython.GkylArray | None: + """The native ``GkylArray`` when gkyl-backed; None otherwise. This is the + handle the modal verbs pass to the Gkeyll kernels.""" + return self._values if isinstance(self._values, gpython.GkylArray) else None + + def get_values(self) -> np.ndarray: + """Values for *reading*: gkyl-backed data yields a read-only NumPy view of + the C buffer (valid while this dataset is alive); numpy-backed data yields + the array itself. Mutation of modal data must go through the kernels.""" + if isinstance(self._values, gpython.GkylArray): + return self._values.view(self.ctx.get("cells")) + return self._values + + def set_values(self, values) -> None: + """Replace stored values and update cell/component metadata.""" + self._values = values + if isinstance(values, gpython.GkylArray): + # Cell layout is not derivable from the flat native array; it comes from + # ctx (set by the reader, and carried through metadata-only copies). + self.ctx["num_comps"] = values.ncomp + else: + self.ctx["cells"] = np.array(values.shape[:-1], dtype=np.int64) + self.ctx["num_comps"] = int(values.shape[-1]) + + values = property(get_values, set_values) + + def __getitem__(self, index): + """Index values using their ordinary NumPy axis order. + + The value layout is ``(*spatial_axes, component)``, so, for example, a + one-dimensional four-component dataset supports ``data[:, 2:4]``. This + deliberately mirrors indexing ``data.values`` rather than treating every + subscript as a component-only selector. A component is selected explicitly with ``data[..., 1]``. + """ + if self._values is None: + raise ValueError("GData values are not loaded; cannot subscript.") + return self.get_values()[index] + + def __setitem__(self, index, value) -> None: + """Assign through NumPy-style indexing on NumPy-backed data. + + Native Gkeyll storage is intentionally read-only from Python; use the + appropriate operation/kernel, or interpolate first, before mutating it. + """ + if self._values is None: + raise ValueError( + "GData values are not loaded; cannot assign by subscript.") + if isinstance(self._values, gpython.GkylArray): + raise ValueError( + "Cannot assign through indexing to native Gkeyll storage; call " + ".interpolate() first to obtain mutable NumPy-backed values.") + self._values[index] = value + + def push(self, grid, values): + """Set values (updating cell/comp ctx) then the grid (updating bounds).""" + self.set_values(values) + self.set_grid(grid) + return self + + # ------------------------------------------------------------- duplication + def clone(self, metadata_only: bool = False) -> "GDataState": + """Deep-copy without re-reading. Builds ``type(self)`` so subclasses + (e.g. the fluent ``GData``) propagate through every verb result. + + Set ``metadata_only=True`` to omit the grid and values from the copy. + """ + new = type(self)(tag=self._tag, label=self._custom_label, ctx=self.ctx) + new.set_label(self._label) + new._file_name = self._file_name + new.color = self.color + if not metadata_only and self._values is not None: + dup = (self._values.clone() if isinstance(self._values, gpython.GkylArray) + else np.array(self._values, copy=True)) + new.push([np.array(g, copy=True) for g in self._grid], dup) + return new + + def _result(self, + grid, + values, + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None, + **ctx_updates): + """The single 'mutate self vs. emit a new dataset' decision point. + + Every verb funnels its computed ``(grid, values)`` through here. Because + ``copy`` uses ``type(self)``, the result is the *same* (sub)class as the + input -- so ``operations`` can be typed on ``GDataState`` yet return a fluent + ``GData`` at runtime. + """ + target = self if inplace else self.clone(metadata_only=True) + target.push(grid, values) + if tag is not None: + target.set_tag(tag) + if label is not None: + target._custom_label = label + if ctx_updates: + target.ctx.update(ctx_updates) + return target + + # ---------------------------------------------------------- operability + @property + def is_interpolated(self) -> bool: + """True when values are safe for element-wise math: data with no DG + structure at all (no ``basis_type`` -- plain point values by + construction), never-modal DG data (``value_form`` is ``nodal``/ + ``quad``), or modal data already run through ``interpolate`` + (``ctx['interpolated']``).""" + if not self.ctx.get("basis_type"): + return True + return (self.ctx.get("value_form", "modal") != "modal" + or self.ctx.get("interpolated", False)) + + def _require_operable(self) -> None: + """Pointwise math is allowed exactly where the data are point values: + the NumPy field domain, or the nodal/quad value forms. Modal + coefficients refuse -- a pointwise operation has no basis-space meaning. + ``value_form`` applies uniformly regardless of ``backend`` -- it is the + one fact for "what do these values mean", set once at load time.""" + if self._values is None: + raise ValueError("GData has no values to operate on.") + if not self.is_interpolated: + raise ValueError( + "Cannot do NumPy math on modal DG coefficients. Convert explicitly: " + ".to_nodal()/.to_quad() (pointwise, stays native), .apply(fn) " + "(pointwise via quadrature, projects back to modal), or .interpolate() " + "(leave for the NumPy field domain).") + + # ----------------------------------------------------- numpy interop (read) + _HANDLED_TYPES = (numbers.Number, np.ndarray, np.generic) + + def __array__(self, dtype=None): + """Expose values so ``np.asarray(data)`` / matplotlib accept the dataset. + + This is a pure *reader* (no ``operations``), so it lives on the container; the + computing operators (``__add__``, ``__array_ufunc__``) live on the fluent + subclass -- see HIERARCHY_3.md. Nodal/quad data expose their point values; + native *modal* data refuses: silently handing out DG coefficients as if + they were point values is a correctness trap.""" + if isinstance(self._values, gpython.GkylArray): + if self.ctx.get("value_form", "modal") != "modal": + return np.asarray(self.get_values(), dtype=dtype) + raise ValueError( + "This dataset holds modal DG coefficients in native Gkeyll storage; " + ".to_nodal()/.to_quad() for point values, or .interpolate() for NumPy." + ) + return np.asarray(self._values, dtype=dtype) + + # -------------------------------------------------------------- reporting + def info(self, index: int = 0, no_header: bool = False) -> str: + """Build and print a summary; optionally omit its descriptive heading.""" + values, num_comps = self.get_values(), self.num_comps + num_dims, num_cells = self.num_dims, self.num_cells + lo, up = self.bounds + out = "" + if not no_header: + lbl = self.get_label() + out += f"{lbl}{' ' if lbl else ''}({self.get_tag()}#{index})\n" + if "time" in self.ctx: + out += f"├─ Time: {self.ctx['time']:e}\n" + if "frame" in self.ctx: + out += f"├─ Frame: {self.ctx['frame']:d}\n" + if self.ctx.get("block") is not None: + out += f"├─ Block: {self.ctx['block']:d} (sim '{self.ctx.get('sim', '')}')\n" + out += f"├─ Number of components: {num_comps:d}\n" + out += f"├─ Number of dimensions: {num_dims:d}\n" + if lo is not None: + out += f"├─ Grid: ({self.get_grid_type()})\n" + for d in range(num_dims): + branch = "└" if d == num_dims - 1 else "├" + out += (f"│ {branch}─ Dim {d}: Num. cells: {int(num_cells[d]):d}; " + f"Lower: {lo[d]:e}; Upper: {up[d]:e}\n") + if values is not None: + vmax = np.nanmax(values) + vmin = np.nanmin(values) + max_idx = np.unravel_index(np.nanargmax(values), values.shape) + min_idx = np.unravel_index(np.nanargmin(values), values.shape) + max_pos = tuple(int(i) for i in max_idx[:num_dims]) + min_pos = tuple(int(i) for i in min_idx[:num_dims]) + out += f"├─ Maximum: {vmax:e} at {max_pos}" + out += f" component {int(max_idx[-1]):d}\n" if num_comps > 1 else "\n" + out += f"├─ Minimum: {vmin:e} at {min_pos}" + out += f" component {int(min_idx[-1]):d}\n" if num_comps > 1 else "\n" + if self.ctx.get("basis_type"): + form = self.ctx.get("value_form", "modal") + if self.ctx.get("interpolated"): + form = "interpolated" + elif form == "quad" and self.ctx.get("num_quad"): + form = f"quad, num_quad={self.ctx['num_quad']}" + out += f"├─ DG: {self.ctx['basis_type']} p{self.ctx.get('poly_order', '?')} ({form})\n" + if "changeset" in self.ctx or "builddate" in self.ctx: + out += "├─ Created with Gkeyll:\n" + if "changeset" in self.ctx: + out += f"│ ├─ Changeset: {self.ctx['changeset']}\n" + if "builddate" in self.ctx: + out += f"│ └─ Build Date: {self.ctx['builddate']}\n" + if "geometry_type" in self.ctx or "geqdsk_sign_convention" in self.ctx: + out += "├─ Geometry info:\n" + if "geometry_type" in self.ctx: + out += f"│ ├─ Type: {self.ctx['geometry_type']}\n" + if "geqdsk_sign_convention" in self.ctx: + out += f"│ ├─ GEQDSK sign convention: {self.ctx['geqdsk_sign_convention']:d}\n" + if any(k in self.ctx for k in ("mass", "charge", "gas_gamma", "vdim")): + out += "├─ Species properties:\n" + if "mass" in self.ctx: + out += f"│ ├─ Mass: {self.ctx['mass']:e}\n" + if "charge" in self.ctx: + out += f"│ ├─ Charge: {self.ctx['charge']:e}\n" + if "gas_gamma" in self.ctx: + out += f"│ ├─ Adiabatic index: {self.ctx['gas_gamma']:e}\n" + if "vdim" in self.ctx: + out += f"│ ├─ Velocity dimensions: {self.ctx['vdim']:d}\n" + for key, val in self.ctx.items(): + if key not in self._INFO_HANDLED_CTX_KEYS: + out += f"├─ {key}: {val}\n" + out += "└─ File: " + (self._file_name or "") + "\n" + print(out) + return out + + # Keys already rendered by a dedicated branch above; anything else in ctx + # is file/reader-native metadata (e.g. a .gkyl file's msgpack meta) that + # still deserves to surface, so it falls through to the generic dump. + _INFO_HANDLED_CTX_KEYS = frozenset({ + "time", + "frame", + "sim", + "block", + "quantity", + "lower", + "upper", + "cells", + "grid_type", + "poly_order", + "basis_type", + "num_comps", + "value_form", + "num_quad", + "interpolated", + "changeset", + "builddate", + "geometry_type", + "geqdsk_sign_convention", + "mass", + "charge", + "gas_gamma", + "vdim", + }) + + # --------------------------------------------------------------- summary + def _summary(self) -> str: + if self._values is None: + return f"<{type(self).__name__} empty | tag '{self._tag}'>" + cells = tuple(int(c) for c in self.get_num_cells()) + parts = [f"<{type(self).__name__} {cells}", f"{self.num_comps:d} comp"] + lo, up = self.bounds + if lo is not None: + parts.append(" ".join(f"[{lo[d]:g},{up[d]:g}]" + for d in range(self.num_dims))) + value_form = self.ctx.get("value_form", "modal") + if self.ctx.get("basis_type"): + dg = str(self.ctx["basis_type"]) + if self.ctx.get("poly_order") is not None: + dg += f" p{self.ctx['poly_order']}" + if self.ctx.get("interpolated"): + dg += " interpolate" + elif value_form == "modal": + dg += " modal" + parts.append(dg) + if self.backend == "gkyl": + parts.append("gkyl-native" if value_form == + "modal" else f"gkyl-native ({value_form})") + parts.append(f"tag '{self._tag}'") + return " | ".join(parts) + ">" + + def __repr__(self) -> str: + return self._summary() + + def __str__(self) -> str: + if self._values is None: + return self._summary() + return (f"{self._summary()}\n" + f"{np.array2string(self.get_values(), threshold=20, edgeitems=2)}") diff --git a/src/postgkyl/gdatastate/gdatastategroup.py b/src/postgkyl/gdatastate/gdatastategroup.py new file mode 100644 index 00000000..d0dfeecb --- /dev/null +++ b/src/postgkyl/gdatastate/gdatastategroup.py @@ -0,0 +1,109 @@ +"""``GDataStateGroup`` -- an ordered, verb-less collection of datasets. + +The container counterpart of :class:`~postgkyl.gdatastate.gdatastate.GDataState`: a group +holds several datasets and offers only *state*-reading operations -- +construction/flattening, the sequence protocol, combining, and a summary +``repr``. Like ``GDataState`` it knows nothing about verbs: no ``operations`` call, +no matplotlib, ever, and it imports only downward (``collection``/``state``, +both in ``gdatastate``). The fluent group that *broadcasts* verbs over its members +(``interpolate``, ``select``, ``plot``, ``info``, ...) is layer 10's job, one layer up +-- exactly the way :class:`postgkyl.gdata.gdata.GData` adds verb methods on top +of ``GDataState`` without ``gdatastate`` ever importing ``api`` (see +:class:`postgkyl.gdata.gdatagroup.GDataGroup`). +""" + +from __future__ import annotations + +from postgkyl.gdatastate.collection import flatten_datasets +from postgkyl.gdatastate.gdatastate import GDataState + + +class GDataStateGroup: + """An ordered collection of ``GDataState`` (or subclass) members. + + Flattens nested lists/tuples/groups of datasets into one ordered sequence + and exposes the sequence protocol (``len``, iteration, indexing/slicing), + combining (``with_``/``&``), and a summary ``repr``. Members keep their own + identity; a group owns no data beyond the ordering. + """ + + def __init__(self, datasets=()): + """Build a group, flattening nested containers of datasets. + + Args: + datasets: GDataState | Iterable + A single dataset, or an (optionally nested) iterable of datasets + and/or other groups. Everything is flattened into one ordered list + via :func:`postgkyl.gdatastate.collection.flatten_datasets`. Defaults to an + empty group. + + Raises: + TypeError: If, after flattening, any member is not a ``GDataState``. + """ + members = flatten_datasets(datasets) if datasets else [] + for member in members: + if not isinstance(member, GDataState): + raise TypeError( + f"Expected a GDataState (or iterable of them), got {type(member)!r}." + ) + self._datasets: list = members + + # ------------------------------------------------------------ sequence + def __iter__(self): + return iter(self._datasets) + + def __len__(self) -> int: + return len(self._datasets) + + def __getitem__(self, index): + """Index or slice the group. + + Args: + index: int | slice + An integer position selects and returns a single member; a + ``slice`` selects a contiguous range. + + Returns: + GDataState | GDataStateGroup: The single member at an integer + ``index``, or a new ``GDataStateGroup`` wrapping the selected members + for a ``slice``. + """ + result = self._datasets[index] + return GDataStateGroup(result) if isinstance(index, slice) else result + + @property + def datasets(self) -> list: + """The members as a plain list. + + Returns a defensive (shallow) copy so callers can mutate the returned + list without affecting this group. + + Returns: + list: A new ``list`` of the members, in order. + """ + return list(self._datasets) + + # ------------------------------------------------------------ combining + def with_(self, *others) -> "GDataStateGroup": + """Return a new group with additional datasets appended. + + Does not mutate this group. ``__and__`` is an alias for this method, so + ``a & b`` is equivalent to ``a.with_(b)``. + + Args: + *others: GDataState | Iterable + Additional members to append. Each may be a single dataset, a + ``GDataStateGroup``, or an (optionally nested) iterable of them; all + are flattened into the resulting group. + + Returns: + GDataStateGroup: A new group containing this group's members followed + by the flattened ``others``. + """ + return GDataStateGroup(self._datasets + list(others)) + + __and__ = with_ + + # ---------------------------------------------------------------- repr + def __repr__(self) -> str: + return f"<{type(self).__name__} [{len(self._datasets):d} datasets]>" diff --git a/src/postgkyl/gdatastate/guards.py b/src/postgkyl/gdatastate/guards.py new file mode 100644 index 00000000..4cea688e --- /dev/null +++ b/src/postgkyl/gdatastate/guards.py @@ -0,0 +1,35 @@ +"""The shared field-domain guard used by field-only verbs and diagnostics. + +Centralizes the check-and-raise boilerplate that was independently retyped +across several ``operations`` physics verbs (moved to ``diagnostics`` by layer 10): +each caller keeps its own ``reason`` clause (why *this* function's math has +no meaning on raw modal coefficients), but the check itself -- +``backend == "gkyl"`` -> raise with the standard ".interpolate() first" message +shape -- has one home. This is a state-invariant helper, not a verb, so it +lives on ``gdatastate`` (which stays verb-less) rather than ``operations``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .gdatastate import GDataState + + +def require_field_domain(data: "GDataState", who: str, reason: str) -> None: + """Raise if ``data`` is native modal (gkyl-backed) DG coefficients. + + Args: + data: The dataset to check. + who: The verb (or argument) name to name in the error message. + reason: The clause explaining why raw coefficients are unusable here, + e.g. ``"rotating raw DG coefficients would mix basis functions"``. + + Raises: + ValueError: if ``data.backend == "gkyl"``. + """ + if data.backend == "gkyl": + raise ValueError( + f"{who} operates on interpolated (NumPy) values; call .interpolate() " + f"first -- {reason}.") diff --git a/src/postgkyl/gdatastate/materialize.py b/src/postgkyl/gdatastate/materialize.py new file mode 100644 index 00000000..29039d85 --- /dev/null +++ b/src/postgkyl/gdatastate/materialize.py @@ -0,0 +1,35 @@ +"""Materialize point-value state without changing its value form.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl.dg import rep + +if TYPE_CHECKING: + from .gdatastate import GDataState + + +def materialize_point_values(data: "GDataState") -> "GDataState": + """Return a NumPy-backed view of nodal or quadrature point values. + + NumPy-backed data is already materialized. Native modal coefficients have + no unique point-value interpretation and therefore require an explicit + representation choice by the caller. + """ + if data.backend != "gkyl": + return data + value_form = data.ctx.get("value_form", "modal") + if value_form == "modal": + raise ValueError( + "modal DG coefficients are not plottable; choose explicitly: " + ".interpolate() (uniform evaluation mesh), .to_nodal() or .to_quad() " + "(plot at the basis/quadrature points).") + grid, values = rep.materialize(str(data.ctx["basis_type"]), data.num_dims, + int(data.ctx["poly_order"]), + data.native, data.grid, value_form, + data.ctx.get("num_quad")) + return data._result(grid, values) + + +__all__ = ["materialize_point_values"] diff --git a/src/postgkyl/gpython/__init__.py b/src/postgkyl/gpython/__init__.py new file mode 100644 index 00000000..b119d245 --- /dev/null +++ b/src/postgkyl/gpython/__init__.py @@ -0,0 +1,43 @@ +"""``gpython/`` -- the foreign floor: the compiled bridge to Gkeyll. + +A bottom leaf (imports nothing internal). This package is the **only** place +in postgkyl that touches the foreign world, and it does so through a compiled +contract (GKEYLL_C_SHIM.md) rather than runtime declarations: + +- ``csrc/`` ``_gpythonmodule.c`` -- the CPython extension over + ``gkyl_gpython.h``; the gpython shim itself lives in the + gkeyll repo (``core/zero/{gkyl_gpython.h, gpython.c}``, + compiled into ``libg0core.so`` by Gkeyll's own build) +- ``_gpython`` the built extension module -- opaque handles in, ndarrays out +- ``_lib`` loads ``_gpython`` + the ``GPYTHON_API_VERSION`` handshake; + ``available()`` is the single capability switch; + ``build_info()`` reads the generated ``_build_info`` (the + vendored Gkeyll commit + build date, written by + ``scripts/build_gpython.sh``) for ``pgkyl --version`` +- ``array`` :class:`GkylArray` -- Python owner of a native ``gkyl_array`` +- ``basis`` cached Gkeyll basis objects + interpolation/nodal/quad matrices + built by evaluating Gkeyll's own basis through the shim +- ``rio`` file loading through ``gkyl_array_rio`` +- ``kernels`` weak multiply/divide/inverse, coefficient lin-combs, reduce, + integrate + +Representation changes (modal · nodal · quad) are orchestration over this +floor's public functions, not floor primitives themselves -- they live in +``dg/rep.py`` (see CLAUDE.md's "Engine layers" section). + +No struct layout, signature, or calling convention exists in Python: the C +compiler checks all of it against the real ``gkyl_*.h`` headers when the shim +builds, so Gkeyll API drift fails the build instead of corrupting data. + +If the extension is missing, importing still succeeds; ``available()`` +returns False and every entry point raises with build guidance. +""" + +from ._lib import available, build_info, lib_path, require +from .array import GkylArray +from . import basis, kernels, rio + +__all__ = [ + "available", "build_info", "lib_path", "require", "GkylArray", "basis", + "kernels", "rio" +] diff --git a/src/postgkyl/gpython/_lib.py b/src/postgkyl/gpython/_lib.py new file mode 100644 index 00000000..228ba171 --- /dev/null +++ b/src/postgkyl/gpython/_lib.py @@ -0,0 +1,82 @@ +"""Load the compiled ``_gpython`` extension -- the single capability switch. + +The foreign floor is the CPython extension ``postgkyl.gpython._gpython``, built by +``scripts/build_gpython.sh`` against ``gkyl_gpython.h`` -- the gpython shim, which lives +in the gkeyll repo (``core/zero/gpython.c``) and is compiled INTO +``libg0core.so`` by Gkeyll's own build (GKEYLL_C_SHIM.md). There are no +runtime signature declarations and no struct mirrors here: the contract is +enforced by the C compiler at the producer. The one runtime check left is +the ``GPYTHON_API_VERSION`` handshake, which catches a stale ``_gpython.so`` +paired with a newer shim header (or vice versa). + +If the extension is missing, :func:`available` returns False and +:func:`require` raises with build guidance; importing postgkyl never fails. + +A NumPy ABI mismatch (`_gpython.so` compiled against a different NumPy than +the one installed -- e.g. pip's isolated build environment resolving a +different NumPy than the target environment did) surfaces as a ``ValueError`` +from NumPy's own ``import_array()`` check, not an ``ImportError``; it is +caught alongside the missing-extension case for the same reason -- a broken +bridge must degrade to "unavailable", never take down ``import postgkyl``. +""" + +from __future__ import annotations + +import pathlib + +try: + from . import _gpython as _mod + if _mod.api_version() != _mod.GPYTHON_API_VERSION: + raise ImportError( + f"gpython shim version mismatch: _gpython.so was built for API " + f"{_mod.api_version()}, postgkyl expects {_mod.GPYTHON_API_VERSION}; " + "rebuild with scripts/build_gpython.sh") + _ERROR = None +except (ImportError, ValueError) as exc: + _mod = None + _ERROR = (f"{exc}\nBuild the compiled bridge with scripts/build_gkeyll.sh " + "(or scripts/build_gpython.sh if libg0core.so already exists). " + "A 'numpy.dtype size changed' error means _gpython.so was " + "compiled against a different NumPy than the one installed here " + "-- reinstall with `pip install -e . --no-build-isolation` so " + "the build step and the installed environment use the same " + "NumPy, then rebuild the bridge.") + + +def available() -> bool: + """True when the compiled Gkeyll bridge is loaded (the capability switch).""" + return _mod is not None + + +def require(): + """The ``_gpython`` module, or a RuntimeError explaining how to build it.""" + if _mod is None: + raise RuntimeError(f"postgkyl's Gkeyll bridge is unavailable: {_ERROR}") + return _mod + + +def lib_path() -> pathlib.Path | None: + """Path of the loaded extension (which is rpath-bound to its libg0core).""" + return pathlib.Path(_mod.__file__) if _mod is not None else None + + +def build_info() -> dict[str, str] | None: + """Metadata about the vendored Gkeyll build this bridge was compiled from. + + None when the bridge has never been built (scripts/build_gkeyll.sh never + ran): ``_build_info`` is a generated build artifact, not part of the + source tree (see scripts/build_gpython.sh, .gitignore). + """ + try: + from . import _build_info as _bi + except ImportError: + return None + return { + "gkeyll_commit": _bi.GKEYLL_COMMIT, + "gkeyll_commit_date": _bi.GKEYLL_COMMIT_DATE, + "gkeyll_branch": _bi.GKEYLL_BRANCH, + "postgkyl_build_commit": _bi.POSTGKYL_BUILD_COMMIT, + "build_date": _bi.BUILD_DATE, + "build_cc": _bi.BUILD_CC, + "build_arch_flags": _bi.BUILD_ARCH_FLAGS + } diff --git a/src/postgkyl/gpython/array.py b/src/postgkyl/gpython/array.py new file mode 100644 index 00000000..db93dfc7 --- /dev/null +++ b/src/postgkyl/gpython/array.py @@ -0,0 +1,93 @@ +"""``GkylArray`` -- the Python owner of a native ``gkyl_array``. + +The handle is a ``PyCapsule`` produced by the ``_gpython`` extension; its +destructor releases the C array, and zero-copy constructions pin the backing +NumPy buffer inside the capsule for the lifetime of the C view. Views of the +data take the capsule as their ndarray ``base``, so a view can never outlive +the native memory it aliases. No raw pointer ever reaches Python. +""" + +from __future__ import annotations + +import numpy as np + +from . import _lib + + +class GkylArray: + """Owns one native ``gkyl_array`` (double-precision) via its capsule.""" + + def __init__(self, cap): + self._cap = cap + + # ------------------------------------------------------------ constructors + @classmethod + def alloc(cls, ncomp: int, size: int) -> "GkylArray": + """gkyl-owned zeroed array of ``size`` cells x ``ncomp`` doubles. + + Raises: + ValueError: ``ncomp`` or ``size`` is not positive. Gkeyll's own + allocator asserts on a zero-byte buffer at *release* time (an abort, + not a Python exception) -- refusing here turns a process crash into a + clean, early error. + """ + if ncomp <= 0 or size <= 0: + raise ValueError(f"GkylArray.alloc: ncomp={ncomp} and size={size} " + "must both be positive (Gkeyll cannot allocate a " + "zero-sized array)") + return cls(_lib.require().array_new(ncomp, size)) + + @classmethod + def from_numpy(cls, values: np.ndarray) -> "GkylArray": + """Zero-copy ``gkyl_array`` view of a ``(cells..., ncomp)`` NumPy array. + + The buffer is pinned inside the capsule for the C array's lifetime; data + is made contiguous float64 first (copying only if needed). + + Raises: + ValueError: ``values`` has fewer than 1 dimension, or is empty (see + :meth:`alloc` -- an empty buffer crashes Gkeyll's allocator on + release rather than raising). + """ + buf = np.ascontiguousarray(values, dtype=np.float64) + if buf.ndim < 1: + raise ValueError("GkylArray.from_numpy: need at least a 1-D " + "(…, ncomp) array") + if buf.size == 0: + raise ValueError("GkylArray.from_numpy: array is empty (Gkeyll " + "cannot allocate a zero-sized array)") + return cls(_lib.require().array_from_numpy(buf)) + + def clone(self) -> "GkylArray": + """Deep copy through ``gkyl_array_clone`` (gkyl-owned).""" + return GkylArray(_lib.require().array_clone(self._cap)) + + # ------------------------------------------------------------------ shape + @property + def ncomp(self) -> int: + return int(_lib.require().array_ncomp(self._cap)) + + @property + def size(self) -> int: + return int(_lib.require().array_size(self._cap)) + + # ---------------------------------------------------------------- readout + def view(self, cells=None) -> np.ndarray: + """Read-only NumPy view of the C buffer, shaped ``(*cells, ncomp)``. + + The view's ``base`` chain holds the owning capsule, so + ``dataset.values.copy()`` on a temporary dataset is safe -- the memory + cannot be released while any view is reachable. Mutation must go through + the kernels, never the view. + """ + flat = _lib.require().array_view(self._cap) + if cells is None: + return flat + return flat.reshape(tuple(int(c) for c in cells) + (flat.shape[-1], )) + + def to_numpy(self, cells=None) -> np.ndarray: + """By-value copy out of the C buffer (what the ``interpolate`` bridge returns).""" + return np.array(self.view(cells), copy=True) + + def __repr__(self) -> str: + return f"" diff --git a/src/postgkyl/gpython/basis.py b/src/postgkyl/gpython/basis.py new file mode 100644 index 00000000..f66fdfc7 --- /dev/null +++ b/src/postgkyl/gpython/basis.py @@ -0,0 +1,303 @@ +"""Gkeyll basis objects + evaluation matrices, through the gpython shim. + +``struct gkyl_basis`` carries the basis functions themselves; the shim +dispatches its function pointers in compiled C (``gpython_basis_eval`` & co.), so +the interpolation matrix is assembled by evaluating Gkeyll's own basis at the +interpolation points -- a few hundred calls, cached per basis -- and NumPy +applies it at array speed. The matrices are therefore bit-consistent with the +kernels the simulation used, with zero layout knowledge in Python. + +Interpolation points follow the historical postgkyl convention: ``num_interp`` +subcell centers per cell, ``z_i = -(n-1)/n + 2 i/n`` on [-1, 1], with +multi-dimensional points ordered Fortran-style (dimension 0 fastest) to match +the per-cell scatter in ``dg/interpolate.py``. +""" + +from __future__ import annotations + +import numpy as np + +from . import _lib + + +class Basis: + """A cached Gkeyll basis: opaque handle + the descriptors postgkyl reads.""" + + def __init__(self, cap, ndim: int, poly_order: int, num_basis: int, id: str): + self._cap = cap + self.ndim = ndim + self.poly_order = poly_order + self.num_basis = num_basis + self.id = id + + def __repr__(self) -> str: + return (f"") + + +_basis_cache: dict[tuple, Basis] = {} +_matrix_cache: dict[tuple, np.ndarray] = {} + +# Highest poly_order each basis supports per ndim, mirroring the fixed-size +# `ev[4]` function-pointer tables in gkeyll's +# core/zero/gkyl_cart_modal_{serendip,tensor}_priv.h. Those tables have NO +# runtime bounds checking: gkyl_cart_modal_serendip/tensor assert +# `ndim>0 && ndim<=6` (a process abort on failure, not a Python exception), +# and index poly_order into the 4-slot array with no check at all, so an +# out-of-range poly_order is undefined behavior, not a clean failure. This +# guard must run before every call into the shim; keep it in sync with those +# two headers if Gkeyll ever adds higher-order kernels. +_MAX_POLY_ORDER = { + "serendipity": { + 1: 3, + 2: 3, + 3: 3, + 4: 3, + 5: 2, + 6: 1 + }, + "tensor": { + 1: 3, + 2: 3, + 3: 2, + 4: 2, + 5: 2, + 6: 1 + }, +} + +# hybrid/gkhybrid are fixed-poly_order (=1) bases parameterized by +# (cdim, vdim) rather than (ndim, poly_order) -- see +# gkeyll/core/zero/gkyl_cart_modal_{hybrid,gkhybrid}.c. A .gkyl file only +# records the total ndim, not the cdim/vdim split, so this table recovers it +# from the one configuration Gkeyll actually produces for each: PKPM hybrid +# always carries a single parallel-velocity direction (vdim=1, cdim=ndim-1); +# gyrokinetic gkhybrid always carries (vpar, mu) (vdim=2), except the 1x1v +# case, which has no mu direction (vdim=1) -- mirroring the legacy postgkyl +# convention (src_bak/postgkyl/data/{dg.py,computeInterpolationMatrices.py}) +# and matching gkeyll/core/unit/ctest_basis.c's own (cdim, vdim) choices. +# Gkeyll's gkhybrid kernel tables are indexed by ndim alone (poly_order fixed +# at 1), so any (cdim, vdim) pair summing to the same ndim would dispatch to +# the identical compiled basis; this table simply names the one physical +# configuration that split corresponds to. +_HYBRID_CDIM_VDIM = { + "hybrid": { + 2: (1, 1), + 3: (2, 1), + 4: (3, 1) + }, + "gkhybrid": { + 2: (1, 1), + 3: (1, 2), + 4: (2, 2), + 5: (3, 2) + }, +} + + +def get_basis(basis_type: str, ndim: int, poly_order: int) -> Basis: + """A cached, fully-initialized Gkeyll basis object. + + Args: + basis_type: ``"serendipity"``, ``"tensor"``, ``"hybrid"``, or + ``"gkhybrid"`` (case-insensitive). + ndim: number of dimensions. 1..6 for serendipity/tensor; the hybrid + bases only exist for the ``(cdim, vdim)`` combinations Gkeyll actually + generates kernels for -- see :data:`_HYBRID_CDIM_VDIM`. + poly_order: polynomial order for serendipity/tensor (ceiling depends on + ``(basis_type, ndim)``, see :data:`_MAX_POLY_ORDER`); must be ``1`` + for hybrid/gkhybrid, which have no other order. + + Returns: + The cached :class:`Basis` (the same object for repeated requests with + the same arguments). + + Raises: + ValueError: unknown ``basis_type``, or ``(ndim, poly_order)`` outside + what Gkeyll's compiled kernel tables support for it. Checked here + because the C constructors have no such guard themselves (see above). + """ + basis_type = basis_type.lower() + key = (basis_type, ndim, poly_order) + if key in _basis_cache: + return _basis_cache[key] + + if basis_type in _HYBRID_CDIM_VDIM: + if poly_order != 1: + raise ValueError(f"Gkeyll's {basis_type} basis only exists at " + f"poly_order 1, got {poly_order}") + cdim_vdim = _HYBRID_CDIM_VDIM[basis_type].get(ndim) + if cdim_vdim is None: + raise ValueError(f"Gkeyll's {basis_type} basis supports ndim " + f"{sorted(_HYBRID_CDIM_VDIM[basis_type])}, got {ndim}") + cdim, vdim = cdim_vdim + cap = _lib.require().basis_new_hybrid(basis_type, cdim, vdim) + else: + limits = _MAX_POLY_ORDER.get(basis_type) + if limits is None: + raise ValueError( + f"unknown basis_type '{basis_type}'; expected one of " + f"{sorted(set(_MAX_POLY_ORDER) | set(_HYBRID_CDIM_VDIM))}") + max_p = limits.get(ndim) + if max_p is None: + raise ValueError(f"Gkeyll's {basis_type} basis supports ndim 1..6, " + f"got {ndim}") + if not 0 <= poly_order <= max_p: + raise ValueError(f"Gkeyll's {basis_type} basis in {ndim}D supports " + f"poly_order 0..{max_p}, got {poly_order}") + cap = _lib.require().basis_new(basis_type, ndim, poly_order) + + nd, p, nb, bid = _lib.require().basis_info(cap) + _basis_cache[key] = Basis(cap, nd, p, nb, bid) + return _basis_cache[key] + + +def num_basis(basis_type: str, ndim: int, poly_order: int) -> int: + """Number of DG basis functions, straight from Gkeyll.""" + return get_basis(basis_type, ndim, poly_order).num_basis + + +def cdim_vdim(basis_type: str, ndim: int) -> tuple[int, int]: + """``(cdim, vdim)`` for a basis: the hybrid/gkhybrid split from + :data:`_HYBRID_CDIM_VDIM`, or ``(ndim, 0)`` for a plain configuration-space + basis (serendipity/tensor have no velocity-space concept). + + Raises: + ValueError: ``basis_type`` is hybrid/gkhybrid and ``ndim`` is not one of + the ``(cdim, vdim)`` combinations Gkeyll generates kernels for. + """ + basis_type = basis_type.lower() + if basis_type in _HYBRID_CDIM_VDIM: + cdim_vdim_ = _HYBRID_CDIM_VDIM[basis_type].get(ndim) + if cdim_vdim_ is None: + raise ValueError(f"Gkeyll's {basis_type} basis supports ndim " + f"{sorted(_HYBRID_CDIM_VDIM[basis_type])}, got {ndim}") + return cdim_vdim_ + return (ndim, 0) + + +def interpolation_points_1d(num_interp: int) -> np.ndarray: + """Subcell-center evaluation points on [-1, 1] (legacy postgkyl convention).""" + n = num_interp + return np.array([-(n - 1.0) / n + 2.0 * i / n for i in range(n)]) + + +def tensor_points(pts_1d: np.ndarray, ndim: int) -> np.ndarray: + """``(len(pts_1d)**ndim, ndim)`` tensor-product point set, dimension 0 + fastest (Fortran multi-index order -- the convention every consumer uses).""" + n = len(pts_1d) + shape = (n, ) * ndim + out = np.empty((n**ndim, ndim)) + for i in range(n**ndim): + idx = np.unravel_index(i, shape, order="F") + out[i, :] = [pts_1d[idx[d]] for d in range(ndim)] + return out + + +def eval_matrix(basis_type: str, ndim: int, poly_order: int, + points: np.ndarray) -> np.ndarray: + """``(npts, num_basis)`` matrix ``M[i, j] = b_j(z_i)`` at arbitrary points + in the reference cell [-1, 1]^ndim -- built by evaluating Gkeyll's own basis + through the shim. The workhorse behind every value_form change *and* + the plotting bridge.""" + g0 = _lib.require() + basis = get_basis(basis_type, ndim, poly_order) + points = np.atleast_2d(np.asarray(points, dtype=np.float64)) + mat = np.empty((points.shape[0], basis.num_basis)) + for i, pt in enumerate(points): + mat[i, :] = g0.basis_eval(basis._cap, pt) + return mat + + +def _cached(key, build): + if key not in _matrix_cache: + mat = build() + mat.flags.writeable = False + _matrix_cache[key] = mat + return _matrix_cache[key] + + +def interpolation_matrix(basis_type: str, ndim: int, poly_order: int, + num_interp: int) -> np.ndarray: + """Evaluation matrix at ``num_interp`` subcell centers per dimension. + + Row ``i`` corresponds to the point with multi-index + ``np.unravel_index(i, [num_interp]*ndim, order="F")`` -- dimension 0 fastest, + matching the consumer in ``dg/interpolate.py``. + """ + return _cached(("interpolation", basis_type, ndim, poly_order, num_interp), + lambda: eval_matrix( + basis_type, ndim, poly_order, + tensor_points(interpolation_points_1d(num_interp), ndim))) + + +# ------------------------------------------------- nodal <-> modal (exact) +def node_coords(basis_type: str, ndim: int, poly_order: int) -> np.ndarray: + """``(num_basis, ndim)`` node coordinates from the basis ``node_list``.""" + basis = get_basis(basis_type, ndim, poly_order) + return _lib.require().basis_node_list(basis._cap) + + +def nodal_to_modal_matrix(basis_type: str, ndim: int, + poly_order: int) -> np.ndarray: + """Exact N×N change of basis, from Gkeyll's ``nodal_to_modal`` + (columns = images of the nodal unit vectors).""" + + def build(): + g0 = _lib.require() + basis = get_basis(basis_type, ndim, poly_order) + nb = basis.num_basis + mat = np.empty((nb, nb)) + for j in range(nb): + fin = np.zeros(nb) + fin[j] = 1.0 + mat[:, j] = g0.basis_nodal_to_modal(basis._cap, fin) + return mat + + return _cached(("n2m", basis_type, ndim, poly_order), build) + + +def modal_to_nodal_matrix(basis_type: str, ndim: int, + poly_order: int) -> np.ndarray: + """Evaluation at the basis nodes -- the exact inverse of ``nodal_to_modal``.""" + return _cached(("m2n", basis_type, ndim, poly_order), lambda: eval_matrix( + basis_type, ndim, poly_order, node_coords(basis_type, ndim, poly_order))) + + +# ------------------------------------------- quadrature <-> modal (projection) +def gauss_quad(ndim: int, num_quad: int): + """Tensor-product Gauss–Legendre rule on [-1, 1]^ndim: + ``(points (nq**ndim, ndim), weights (nq**ndim,))``, dimension 0 fastest.""" + p1, w1 = np.polynomial.legendre.leggauss(num_quad) + pts = tensor_points(p1, ndim) + shape = (num_quad, ) * ndim + w = np.empty(num_quad**ndim) + for i in range(w.size): + idx = np.unravel_index(i, shape, order="F") + w[i] = np.prod([w1[idx[d]] for d in range(ndim)]) + return pts, w + + +def modal_to_quad_matrix(basis_type: str, ndim: int, poly_order: int, + num_quad: int) -> np.ndarray: + """``(nq**ndim, num_basis)`` -- evaluate the expansion at the Gauss points.""" + return _cached(("m2q", basis_type, ndim, poly_order, num_quad), + lambda: eval_matrix(basis_type, ndim, poly_order, + gauss_quad(ndim, num_quad)[0])) + + +def quad_to_modal_matrix(basis_type: str, ndim: int, poly_order: int, + num_quad: int) -> np.ndarray: + """``(num_basis, nq**ndim)`` quadrature projection ``c_j = sum_i w_i b_j(z_i) f_i``. + + Exact whenever the integrand ``f·b_j`` has degree ≤ 2·num_quad−1 (the bases + are orthonormal on the reference cell, so no mass-matrix solve is needed). + ``quad_to_modal @ modal_to_quad == I`` for ``num_quad >= p+1``. + """ + + def build(): + pts, w = gauss_quad(ndim, num_quad) + B = eval_matrix(basis_type, ndim, poly_order, pts) + return B.T * w # (N, npts): rows b_j(z_i), scaled by the weights + + return _cached(("q2m", basis_type, ndim, poly_order, num_quad), build) diff --git a/src/postgkyl/gpython/csrc/_gpythonmodule.c b/src/postgkyl/gpython/csrc/_gpythonmodule.c new file mode 100644 index 00000000..b907af11 --- /dev/null +++ b/src/postgkyl/gpython/csrc/_gpythonmodule.c @@ -0,0 +1,1086 @@ +/* _gpythonmodule.c -- the CPython extension over gkyl_gpython.h + * (GKEYLL_C_SHIM.md). + * + * Knows Python objects, NumPy arrays, and the gpython contract -- and nothing + * else about Gkeyll: gkyl_gpython.h (the gpython shim, which lives in the + * gkeyll repo and is compiled into libg0core.so) exposes only opaque handles, + * scalars, and buffers, so no layout or calling convention exists on this + * side of the wall. + * + * Ownership model: native handles live in PyCapsules whose destructors + * release the C object; a capsule's context slot optionally pins a Python + * buffer the C array aliases (zero-copy construction). NumPy views of + * array data take the capsule as their base, so a view keeps the native + * memory alive for as long as the view itself is reachable. + */ +#define PY_SSIZE_T_CLEAN +#include + +#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION +/* Pin the NumPy ABI this extension targets to postgkyl's own declared floor + * (numpy>=2.2.6 in pyproject.toml). Without this, NumPy's headers default + * PyArray_Descr et al. to whatever the *compiling* NumPy's minor version + * happens to be, so a `_gpython.so` built against e.g. NumPy 2.5 headers can + * fail NumPy's own import_array() ABI check ("numpy.dtype size changed") + * against a NumPy 2.2 runtime -- a real hazard here because pip's isolated + * build environment resolves build-system.requires' numpy independently + * from the numpy actually installed into the target environment. + * + * This is a best-effort backstop, NOT a substitute for building against the + * actual runtime NumPy: a build/runtime version skew has been reproduced to + * crash outright (segfault / heap corruption inside Gkeyll's C code, not a + * clean import_array() failure) even with this pin in place -- see + * scripts/build_gpython.sh and README.md's "--no-build-isolation" install + * instructions, which are the real fix. + */ +#define NPY_TARGET_VERSION NPY_2_2_API_VERSION +#include + +#include + +#include + +static const char ARRAY_CAP[] = "gpython_array"; +static const char BASIS_CAP[] = "gpython_basis"; + +/* ------------------------------------------------------------ capsules */ +static void +array_capsule_destroy(PyObject *cap) +{ + gpython_array *a = PyCapsule_GetPointer(cap, ARRAY_CAP); + if (a) + gpython_array_release(a); + Py_XDECREF((PyObject *)PyCapsule_GetContext(cap)); +} + +static void +basis_capsule_destroy(PyObject *cap) +{ + gpython_basis *b = PyCapsule_GetPointer(cap, BASIS_CAP); + if (b) + gpython_basis_release(b); +} + +static gpython_array * +array_arg(PyObject *cap) +{ + return (gpython_array *)PyCapsule_GetPointer(cap, ARRAY_CAP); +} + +static gpython_basis * +basis_arg(PyObject *cap) +{ + return (gpython_basis *)PyCapsule_GetPointer(cap, BASIS_CAP); +} + +static PyObject * +wrap_array(gpython_array *a) +{ + if (!a) { + PyErr_SetString(PyExc_MemoryError, "received NULL gpython_array"); + return NULL; + } + PyObject *cap = PyCapsule_New(a, ARRAY_CAP, array_capsule_destroy); + if (!cap) + gpython_array_release(a); + return cap; +} + +/* --------------------------------------------------------------- misc */ +static PyObject * +py_api_version(PyObject *self, PyObject *noargs) +{ + return PyLong_FromLong(gpython_api_version()); +} + +/* -------------------------------------------------------------- arrays */ +static PyObject * +py_array_new(PyObject *self, PyObject *args) +{ + Py_ssize_t ncomp, size; + if (!PyArg_ParseTuple(args, "nn", &ncomp, &size)) + return NULL; + return wrap_array(gpython_array_new((size_t)ncomp, (size_t)size)); +} + +static PyObject * +py_array_from_numpy(PyObject *self, PyObject *args) +{ + PyObject *obj; + if (!PyArg_ParseTuple(args, "O", &obj)) + return NULL; + /* A C-contiguous float64 array; copies only if the input is not already + * one. The result is pinned in the capsule context: the C array is a + * zero-copy view of exactly this buffer. */ + PyArrayObject *buf = + (PyArrayObject *)PyArray_FROM_OTF(obj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + if (!buf) + return NULL; + if (PyArray_NDIM(buf) < 1) { + Py_DECREF(buf); + PyErr_SetString(PyExc_ValueError, "need at least a 1-D (…, ncomp) array"); + return NULL; + } + npy_intp ncomp = PyArray_DIM(buf, PyArray_NDIM(buf) - 1); + npy_intp size = PyArray_SIZE(buf) / (ncomp ? ncomp : 1); + gpython_array *a = + gpython_array_from_buff((size_t)ncomp, (size_t)size, PyArray_DATA(buf)); + PyObject *cap = wrap_array(a); + if (!cap) { + Py_DECREF(buf); + return NULL; + } + PyCapsule_SetContext(cap, buf); /* pin: released by the capsule dtor */ + return cap; +} + +static PyObject * +py_array_clone(PyObject *self, PyObject *args) +{ + PyObject *cap; + if (!PyArg_ParseTuple(args, "O", &cap)) + return NULL; + gpython_array *a = array_arg(cap); + if (!a) + return NULL; + return wrap_array(gpython_array_clone(a)); +} + +static PyObject * +py_array_ncomp(PyObject *self, PyObject *args) +{ + PyObject *cap; + if (!PyArg_ParseTuple(args, "O", &cap)) + return NULL; + gpython_array *a = array_arg(cap); + if (!a) + return NULL; + return PyLong_FromSize_t(gpython_array_ncomp(a)); +} + +static PyObject * +py_array_size(PyObject *self, PyObject *args) +{ + PyObject *cap; + if (!PyArg_ParseTuple(args, "O", &cap)) + return NULL; + gpython_array *a = array_arg(cap); + if (!a) + return NULL; + return PyLong_FromSize_t(gpython_array_size(a)); +} + +static PyObject * +py_array_view(PyObject *self, PyObject *args) +{ + PyObject *cap; + if (!PyArg_ParseTuple(args, "O", &cap)) + return NULL; + gpython_array *a = array_arg(cap); + if (!a) + return NULL; + npy_intp dims[2] = {(npy_intp)gpython_array_size(a), + (npy_intp)gpython_array_ncomp(a)}; + PyObject *view = + PyArray_SimpleNewFromData(2, dims, NPY_DOUBLE, gpython_array_data(a)); + if (!view) + return NULL; + Py_INCREF(cap); /* base steals this reference */ + if (PyArray_SetBaseObject((PyArrayObject *)view, cap) < 0) { + Py_DECREF(view); + return NULL; + } + PyArray_CLEARFLAGS((PyArrayObject *)view, NPY_ARRAY_WRITEABLE); + return view; +} + +/* ------------------------------------------------------------- file I/O */ +static PyObject * +grid_tuple(int ndim, const double *lower, const double *upper, const int *cells) +{ + npy_intp n = ndim; + PyObject *lo = PyArray_SimpleNew(1, &n, NPY_DOUBLE); + PyObject *up = PyArray_SimpleNew(1, &n, NPY_DOUBLE); + PyObject *nc = PyArray_SimpleNew(1, &n, NPY_INT64); + if (!lo || !up || !nc) { + Py_XDECREF(lo); + Py_XDECREF(up); + Py_XDECREF(nc); + return NULL; + } + for (int d = 0; d < ndim; ++d) { + ((double *)PyArray_DATA((PyArrayObject *)lo))[d] = lower[d]; + ((double *)PyArray_DATA((PyArrayObject *)up))[d] = upper[d]; + ((npy_int64 *)PyArray_DATA((PyArrayObject *)nc))[d] = cells[d]; + } + return Py_BuildValue("(iNNN)", ndim, lo, up, nc); +} + +static PyObject * +py_file_type(PyObject *self, PyObject *args) +{ + const char *fname; + if (!PyArg_ParseTuple(args, "s", &fname)) + return NULL; + return PyLong_FromLong(gpython_file_type(fname)); +} + +static PyObject * +py_read_header(PyObject *self, PyObject *args) +{ + const char *fname; + if (!PyArg_ParseTuple(args, "s", &fname)) + return NULL; + int ndim, file_type, cells[gpython_MAX_DIM]; + double lower[gpython_MAX_DIM], upper[gpython_MAX_DIM]; + size_t esznc, tot_cells, meta_sz; + char *meta; + int status = + gpython_read_header(fname, &ndim, lower, upper, cells, &file_type, &esznc, + &tot_cells, &meta, &meta_sz); + if (status != 0) { + PyErr_Format(PyExc_OSError, "'%s': %s", fname, gpython_status_msg(status)); + return NULL; + } + PyObject *grid = grid_tuple(ndim, lower, upper, cells); + PyObject *meta_bytes = + PyBytes_FromStringAndSize(meta ? meta : "", (Py_ssize_t)meta_sz); + gpython_meta_release(meta); + if (!grid || !meta_bytes) { + Py_XDECREF(grid); + Py_XDECREF(meta_bytes); + return NULL; + } + return Py_BuildValue("(NiNnn)", grid, file_type, meta_bytes, + (Py_ssize_t)esznc, (Py_ssize_t)tot_cells); +} + +static PyObject * +py_read_field(PyObject *self, PyObject *args) +{ + const char *fname; + if (!PyArg_ParseTuple(args, "s", &fname)) + return NULL; + int ndim, cells[gpython_MAX_DIM]; + double lower[gpython_MAX_DIM], upper[gpython_MAX_DIM]; + gpython_array *a = gpython_read_field(fname, &ndim, lower, upper, cells); + if (!a) { + PyErr_Format(PyExc_OSError, "'%s': gpython_read_field failed", fname); + return NULL; + } + PyObject *grid = grid_tuple(ndim, lower, upper, cells); + PyObject *cap = wrap_array(a); + if (!grid || !cap) { + Py_XDECREF(grid); + Py_XDECREF(cap); + return NULL; + } + return Py_BuildValue("(NN)", grid, cap); +} + +/* --------------------------------------------------------------- basis */ +static PyObject * +py_basis_new(PyObject *self, PyObject *args) +{ + const char *type; + int ndim, poly_order; + if (!PyArg_ParseTuple(args, "sii", &type, &ndim, &poly_order)) + return NULL; + gpython_basis *b = gpython_basis_new(type, ndim, poly_order); + if (!b) { + PyErr_Format(PyExc_NotImplementedError, + "basis '%s' is not wired through the Gkeyll shim", type); + return NULL; + } + return PyCapsule_New(b, BASIS_CAP, basis_capsule_destroy); +} + +static PyObject * +py_basis_new_hybrid(PyObject *self, PyObject *args) +{ + const char *type; + int cdim, vdim; + if (!PyArg_ParseTuple(args, "sii", &type, &cdim, &vdim)) + return NULL; + gpython_basis *b = gpython_basis_new_hybrid(type, cdim, vdim); + if (!b) { + PyErr_Format(PyExc_NotImplementedError, + "basis '%s' is not wired through the Gkeyll shim", type); + return NULL; + } + return PyCapsule_New(b, BASIS_CAP, basis_capsule_destroy); +} + +static PyObject * +py_basis_info(PyObject *self, PyObject *args) +{ + PyObject *cap; + if (!PyArg_ParseTuple(args, "O", &cap)) + return NULL; + gpython_basis *b = basis_arg(cap); + if (!b) + return NULL; + return Py_BuildValue("(iiis)", gpython_basis_ndim(b), + gpython_basis_poly_order(b), gpython_basis_num_basis(b), + gpython_basis_id(b)); +} + +static PyObject * +py_basis_eval(PyObject *self, PyObject *args) +{ + PyObject *cap, *zobj; + if (!PyArg_ParseTuple(args, "OO", &cap, &zobj)) + return NULL; + gpython_basis *b = basis_arg(cap); + if (!b) + return NULL; + PyArrayObject *z = + (PyArrayObject *)PyArray_FROM_OTF(zobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + if (!z) + return NULL; + if (PyArray_SIZE(z) < gpython_basis_ndim(b)) { + Py_DECREF(z); + PyErr_SetString(PyExc_ValueError, "point has fewer entries than ndim"); + return NULL; + } + npy_intp nb = gpython_basis_num_basis(b); + PyObject *out = PyArray_SimpleNew(1, &nb, NPY_DOUBLE); + if (!out) { + Py_DECREF(z); + return NULL; + } + gpython_basis_eval(b, PyArray_DATA(z), PyArray_DATA((PyArrayObject *)out)); + Py_DECREF(z); + return out; +} + +static PyObject * +py_basis_node_list(PyObject *self, PyObject *args) +{ + PyObject *cap; + if (!PyArg_ParseTuple(args, "O", &cap)) + return NULL; + gpython_basis *b = basis_arg(cap); + if (!b) + return NULL; + npy_intp dims[2] = {gpython_basis_num_basis(b), gpython_basis_ndim(b)}; + PyObject *out = PyArray_SimpleNew(2, dims, NPY_DOUBLE); + if (!out) + return NULL; + gpython_basis_node_list(b, PyArray_DATA((PyArrayObject *)out)); + return out; +} + +static PyObject * +py_basis_nodal_to_modal(PyObject *self, PyObject *args) +{ + PyObject *cap, *fobj; + if (!PyArg_ParseTuple(args, "OO", &cap, &fobj)) + return NULL; + gpython_basis *b = basis_arg(cap); + if (!b) + return NULL; + PyArrayObject *fin = + (PyArrayObject *)PyArray_FROM_OTF(fobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + if (!fin) + return NULL; + npy_intp nb = gpython_basis_num_basis(b); + if (PyArray_SIZE(fin) != nb) { + Py_DECREF(fin); + PyErr_SetString(PyExc_ValueError, "expected num_basis nodal values"); + return NULL; + } + PyObject *out = PyArray_SimpleNew(1, &nb, NPY_DOUBLE); + if (!out) { + Py_DECREF(fin); + return NULL; + } + gpython_basis_nodal_to_modal(b, PyArray_DATA(fin), + PyArray_DATA((PyArrayObject *)out)); + Py_DECREF(fin); + return out; +} + +/* ------------------------------------------------------ weak DG algebra */ +typedef int (*binop_fn)(const gpython_basis *, gpython_array *, + const gpython_array *, const gpython_array *); + +static PyObject * +binop(PyObject *args, binop_fn fn, const char *name) +{ + PyObject *bcap, *ocap, *acap, *bcap2; + if (!PyArg_ParseTuple(args, "OOOO", &bcap, &ocap, &acap, &bcap2)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *out = array_arg(ocap), *a1 = array_arg(acap), + *a2 = array_arg(bcap2); + if (!b || !out || !a1 || !a2) + return NULL; + if (fn(b, out, a1, a2) != 0) { + PyErr_Format(PyExc_ValueError, + "%s: operand shapes incompatible with " + "the basis (ncomp must be a multiple of num_basis and match)", + name); + return NULL; + } + Py_RETURN_NONE; +} + +static PyObject * +py_dg_mul(PyObject *self, PyObject *args) +{ + return binop(args, gpython_dg_mul, "dg_mul"); +} + +static PyObject * +py_dg_div(PyObject *self, PyObject *args) +{ + return binop(args, gpython_dg_div, "dg_div"); +} + +static PyObject * +py_dg_inv(PyObject *self, PyObject *args) +{ + PyObject *bcap, *ocap, *acap; + if (!PyArg_ParseTuple(args, "OOO", &bcap, &ocap, &acap)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *out = array_arg(ocap), *a1 = array_arg(acap); + if (!b || !out || !a1) + return NULL; + if (gpython_dg_inv(b, out, a1) != 0) { + PyErr_SetString(PyExc_ValueError, + "dg_inv: operand shapes incompatible with the basis"); + return NULL; + } + Py_RETURN_NONE; +} + +static PyObject * +py_dg_mul_conf_phase(PyObject *self, PyObject *args) +{ + PyObject *cbcap, *pbcap, *poutcap, *copcap, *popcap, *ccellsobj, *pcellsobj; + if (!PyArg_ParseTuple(args, "OOOOOOO", &cbcap, &pbcap, &poutcap, &copcap, + &popcap, &ccellsobj, &pcellsobj)) + return NULL; + gpython_basis *cbasis = basis_arg(cbcap), *pbasis = basis_arg(pbcap); + gpython_array *pout = array_arg(poutcap), *cop = array_arg(copcap), + *pop = array_arg(popcap); + if (!cbasis || !pbasis || !pout || !cop || !pop) + return NULL; + PyArrayObject *ccells = (PyArrayObject *)PyArray_FROM_OTF( + ccellsobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + PyArrayObject *pcells = (PyArrayObject *)PyArray_FROM_OTF( + pcellsobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + if (!ccells || !pcells) { + Py_XDECREF(ccells); + Py_XDECREF(pcells); + return NULL; + } + if (PyArray_SIZE(ccells) != gpython_basis_ndim(cbasis) || + PyArray_SIZE(pcells) != gpython_basis_ndim(pbasis)) { + Py_DECREF(ccells); + Py_DECREF(pcells); + PyErr_SetString( + PyExc_ValueError, + "mul_conf_phase: cells arrays must match each basis's ndim"); + return NULL; + } + int status = + gpython_dg_mul_conf_phase(cbasis, pbasis, pout, cop, pop, + PyArray_DATA(ccells), PyArray_DATA(pcells)); + Py_DECREF(ccells); + Py_DECREF(pcells); + if (status != 0) { + PyErr_SetString( + PyExc_ValueError, + "mul_conf_phase: operand shapes incompatible with the bases/cells"); + return NULL; + } + Py_RETURN_NONE; +} + +/* --------------------------------------- linear coefficient ops / reduce */ +static PyObject * +py_array_set(PyObject *self, PyObject *args) +{ + PyObject *ocap, *acap; + double c; + if (!PyArg_ParseTuple(args, "OdO", &ocap, &c, &acap)) + return NULL; + gpython_array *out = array_arg(ocap), *a = array_arg(acap); + if (!out || !a) + return NULL; + gpython_array_set(out, c, a); + Py_RETURN_NONE; +} + +static PyObject * +py_array_accumulate(PyObject *self, PyObject *args) +{ + PyObject *ocap, *acap; + double c; + if (!PyArg_ParseTuple(args, "OdO", &ocap, &c, &acap)) + return NULL; + gpython_array *out = array_arg(ocap), *a = array_arg(acap); + if (!out || !a) + return NULL; + gpython_array_accumulate(out, c, a); + Py_RETURN_NONE; +} + +static PyObject * +py_array_scale(PyObject *self, PyObject *args) +{ + PyObject *acap; + double c; + if (!PyArg_ParseTuple(args, "Od", &acap, &c)) + return NULL; + gpython_array *a = array_arg(acap); + if (!a) + return NULL; + gpython_array_scale(a, c); + Py_RETURN_NONE; +} + +static PyObject * +py_array_shiftc(PyObject *self, PyObject *args) +{ + PyObject *acap; + double val; + unsigned comp; + if (!PyArg_ParseTuple(args, "OdI", &acap, &val, &comp)) + return NULL; + gpython_array *a = array_arg(acap); + if (!a) + return NULL; + gpython_array_shiftc(a, val, comp); + Py_RETURN_NONE; +} + +static PyObject * +py_array_reduce(PyObject *self, PyObject *args) +{ + PyObject *acap; + int op; + if (!PyArg_ParseTuple(args, "Oi", &acap, &op)) + return NULL; + gpython_array *a = array_arg(acap); + if (!a) + return NULL; + if (op < 0 || op > 2) { + PyErr_SetString(PyExc_ValueError, "reduce op must be 0/1/2 (min/max/sum)"); + return NULL; + } + npy_intp n = (npy_intp)gpython_array_ncomp(a); + PyObject *out = PyArray_SimpleNew(1, &n, NPY_DOUBLE); + if (!out) + return NULL; + gpython_array_reduce(PyArray_DATA((PyArrayObject *)out), a, op); + return out; +} + +/* ---------------------------------------------------- field-aware reduce */ +static PyObject * +py_array_dg_reduce(PyObject *self, PyObject *args) +{ + PyObject *bcap, *acap; + int comp, op; + if (!PyArg_ParseTuple(args, "OOii", &bcap, &acap, &comp, &op)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *a = array_arg(acap); + if (!b || !a) + return NULL; + if (op < 0 || op > 2) { + PyErr_SetString(PyExc_ValueError, "reduce op must be 0/1/2 (min/max/sum)"); + return NULL; + } + double out; + if (gpython_array_dg_reduce(&out, b, a, comp, op) != 0) { + PyErr_Format(PyExc_ValueError, + "dg_reduce: component %d out of range for this basis", comp); + return NULL; + } + return PyFloat_FromDouble(out); +} + +/* ------------------------------------------------------------ integrate */ +static PyObject * +py_array_integrate(PyObject *self, PyObject *args) +{ + PyObject *loobj, *upobj, *ncobj, *bcap, *acap; + int nfields, op; + double factor; + if (!PyArg_ParseTuple(args, "OOOOiidO", &loobj, &upobj, &ncobj, &bcap, + &nfields, &op, &factor, &acap)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *a = array_arg(acap); + if (!b || !a) + return NULL; + if (op < 0 || op > 2) { + PyErr_SetString(PyExc_ValueError, + "integrate op must be 0/1/2 (none/abs/sq)"); + return NULL; + } + PyArrayObject *lo = + (PyArrayObject *)PyArray_FROM_OTF(loobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *up = + (PyArrayObject *)PyArray_FROM_OTF(upobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *nc = + (PyArrayObject *)PyArray_FROM_OTF(ncobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + if (!lo || !up || !nc) + goto fail; + int ndim = (int)PyArray_SIZE(lo); + if (ndim < 1 || ndim > gpython_MAX_DIM || PyArray_SIZE(up) != ndim || + PyArray_SIZE(nc) != ndim) { + PyErr_SetString(PyExc_ValueError, "grid arrays must share ndim <= 7"); + goto fail; + } + npy_intp n = nfields; + PyObject *out = PyArray_SimpleNew(1, &n, NPY_DOUBLE); + if (!out) + goto fail; + int status = gpython_array_integrate(ndim, PyArray_DATA(lo), PyArray_DATA(up), + PyArray_DATA(nc), b, nfields, op, factor, + a, PyArray_DATA((PyArrayObject *)out)); + Py_DECREF(lo); + Py_DECREF(up); + Py_DECREF(nc); + if (status != 0) { + Py_DECREF(out); + PyErr_SetString(PyExc_ValueError, + "integrate: grid cells do not cover the array"); + return NULL; + } + return out; +fail: + Py_XDECREF(lo); + Py_XDECREF(up); + Py_XDECREF(nc); + return NULL; +} + +/* ------------------------------------------------------------- average */ +static PyObject * +py_array_average(PyObject *self, PyObject *args) +{ + PyObject *loobj, *upobj, *ncobj, *bcap, *bavgcap, *ncavgobj, *avgdimobj, + *wcap, *acap, *ocap; + if (!PyArg_ParseTuple(args, "OOOOOOOOOO", &loobj, &upobj, &ncobj, &bcap, + &bavgcap, &ncavgobj, &avgdimobj, &wcap, &acap, &ocap)) + return NULL; + gpython_basis *b = basis_arg(bcap), *b_avg = basis_arg(bavgcap); + gpython_array *a = array_arg(acap), *out = array_arg(ocap); + if (!b || !b_avg || !a || !out) + return NULL; + gpython_array *weight = NULL; + if (wcap != Py_None) { + weight = array_arg(wcap); + if (!weight) + return NULL; + } + PyArrayObject *lo = + (PyArrayObject *)PyArray_FROM_OTF(loobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *up = + (PyArrayObject *)PyArray_FROM_OTF(upobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *nc = + (PyArrayObject *)PyArray_FROM_OTF(ncobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + PyArrayObject *ncavg = (PyArrayObject *)PyArray_FROM_OTF(ncavgobj, NPY_INT32, + NPY_ARRAY_IN_ARRAY); + PyArrayObject *avgdim = (PyArrayObject *)PyArray_FROM_OTF( + avgdimobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + if (!lo || !up || !nc || !ncavg || !avgdim) + goto fail; + int ndim = (int)PyArray_SIZE(lo); + int ndim_avg = (int)PyArray_SIZE(ncavg); + if (ndim < 1 || ndim > gpython_MAX_DIM || PyArray_SIZE(up) != ndim || + PyArray_SIZE(nc) != ndim || PyArray_SIZE(avgdim) != ndim) { + PyErr_SetString(PyExc_ValueError, + "average: grid/avg_dim arrays must share ndim <= 7"); + goto fail; + } + if (ndim_avg < 1 || ndim_avg > gpython_MAX_DIM) { + PyErr_SetString(PyExc_ValueError, + "average: cells_avg must have between 1 and 7 entries"); + goto fail; + } + int status = gpython_array_average( + ndim, PyArray_DATA(lo), PyArray_DATA(up), PyArray_DATA(nc), b, b_avg, + ndim_avg, PyArray_DATA(ncavg), PyArray_DATA(avgdim), weight, a, out); + Py_DECREF(lo); + Py_DECREF(up); + Py_DECREF(nc); + Py_DECREF(ncavg); + Py_DECREF(avgdim); + if (status != 0) { + PyErr_SetString(PyExc_ValueError, + "average: operand shapes incompatible with the bases/grid"); + return NULL; + } + Py_RETURN_NONE; +fail: + Py_XDECREF(lo); + Py_XDECREF(up); + Py_XDECREF(nc); + Py_XDECREF(ncavg); + Py_XDECREF(avgdim); + return NULL; +} + +/* --------------------------------------------------------- pow(sqrt) */ +static PyObject * +py_powsqrt(PyObject *self, PyObject *args) +{ + PyObject *bcap, *ncobj, *ocap, *acap; + int num_quad; + double exponent; + if (!PyArg_ParseTuple(args, "OidOOO", &bcap, &num_quad, &exponent, &ncobj, + &ocap, &acap)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *out = array_arg(ocap), *a = array_arg(acap); + if (!b || !out || !a) + return NULL; + PyArrayObject *nc = + (PyArrayObject *)PyArray_FROM_OTF(ncobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + if (!nc) + return NULL; + int ndim = (int)PyArray_SIZE(nc); + if (ndim < 1 || ndim > gpython_MAX_DIM) { + Py_DECREF(nc); + PyErr_SetString(PyExc_ValueError, + "powsqrt: cells must have 1 to 7 entries"); + return NULL; + } + int status = + gpython_powsqrt(b, num_quad, exponent, ndim, PyArray_DATA(nc), out, a); + Py_DECREF(nc); + if (status != 0) { + PyErr_SetString(PyExc_ValueError, + "powsqrt: operand shapes incompatible with the basis/grid"); + return NULL; + } + Py_RETURN_NONE; +} + +/* ------------------------------------------------------------- writing */ +static PyObject * +py_write_field(PyObject *self, PyObject *args) +{ + const char *fname; + PyObject *loobj, *upobj, *ncobj, *metaobj, *acap; + if (!PyArg_ParseTuple(args, "sOOOOO", &fname, &loobj, &upobj, &ncobj, + &metaobj, &acap)) + return NULL; + gpython_array *a = array_arg(acap); + if (!a) + return NULL; + PyArrayObject *lo = + (PyArrayObject *)PyArray_FROM_OTF(loobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *up = + (PyArrayObject *)PyArray_FROM_OTF(upobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *nc = + (PyArrayObject *)PyArray_FROM_OTF(ncobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + if (!lo || !up || !nc) + goto fail; + int ndim = (int)PyArray_SIZE(lo); + if (ndim < 1 || ndim > gpython_MAX_DIM || PyArray_SIZE(up) != ndim || + PyArray_SIZE(nc) != ndim) { + PyErr_SetString(PyExc_ValueError, "grid arrays must share ndim <= 7"); + goto fail; + } + const char *meta = NULL; + Py_ssize_t meta_sz = 0; + if (metaobj != Py_None) { + if (PyBytes_AsStringAndSize(metaobj, (char **)&meta, &meta_sz) < 0) + goto fail; + } + int status = + gpython_write_field(fname, ndim, PyArray_DATA(lo), PyArray_DATA(up), + PyArray_DATA(nc), meta, (size_t)meta_sz, a); + Py_DECREF(lo); + Py_DECREF(up); + Py_DECREF(nc); + if (status == -1) { + PyErr_SetString(PyExc_ValueError, + "write_field: grid cells do not cover the array"); + return NULL; + } + if (status != 0) { + PyErr_Format(PyExc_OSError, "'%s': %s", fname, gpython_status_msg(status)); + return NULL; + } + Py_RETURN_NONE; +fail: + Py_XDECREF(lo); + Py_XDECREF(up); + Py_XDECREF(nc); + return NULL; +} + +/* ------------------------------------------------------- differentiate */ +static PyObject * +py_dg_differentiate(PyObject *self, PyObject *args) +{ + PyObject *bcap, *ocap, *icap; + int dir, diff_order; + double dx; + if (!PyArg_ParseTuple(args, "OiidOO", &bcap, &dir, &diff_order, &dx, &ocap, + &icap)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *out = array_arg(ocap), *in = array_arg(icap); + if (!b || !out || !in) + return NULL; + if (gpython_dg_differentiate(b, dir, diff_order, dx, out, in) != 0) { + PyErr_SetString( + PyExc_ValueError, + "dg_differentiate: operand shapes incompatible with the basis, or " + "dir/diff_order out of range"); + return NULL; + } + Py_RETURN_NONE; +} + +/* ---------------------------------------------------- evaluate-and-project */ +static PyObject * +py_eval_at_coord_proj(PyObject *self, PyObject *args) +{ + PyObject *bcap, *loobj, *upobj, *ncobj, *evaldirsobj, *evalcoordsobj, + *ncavgobj, *icap; + int cdim_do, ndim_tar; + if (!PyArg_ParseTuple(args, "OiOOOOOiOO", &bcap, &cdim_do, &loobj, &upobj, + &ncobj, &evaldirsobj, &evalcoordsobj, &ndim_tar, + &ncavgobj, &icap)) + return NULL; + gpython_basis *b = basis_arg(bcap); + gpython_array *in = array_arg(icap); + if (!b || !in) + return NULL; + PyArrayObject *lo = + (PyArrayObject *)PyArray_FROM_OTF(loobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *up = + (PyArrayObject *)PyArray_FROM_OTF(upobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *nc = + (PyArrayObject *)PyArray_FROM_OTF(ncobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + PyArrayObject *evaldirs = (PyArrayObject *)PyArray_FROM_OTF( + evaldirsobj, NPY_INT32, NPY_ARRAY_IN_ARRAY); + PyArrayObject *evalcoords = (PyArrayObject *)PyArray_FROM_OTF( + evalcoordsobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *ncavg = (PyArrayObject *)PyArray_FROM_OTF(ncavgobj, NPY_INT32, + NPY_ARRAY_IN_ARRAY); + if (!lo || !up || !nc || !evaldirs || !evalcoords || !ncavg) + goto fail; + int ndim = (int)PyArray_SIZE(lo); + int num_eval = (int)PyArray_SIZE(evaldirs); + if (ndim < 1 || ndim > gpython_MAX_DIM || PyArray_SIZE(up) != ndim || + PyArray_SIZE(nc) != ndim) { + PyErr_SetString(PyExc_ValueError, + "eval_at_coord_proj: grid arrays must share ndim <= 7"); + goto fail; + } + if (PyArray_SIZE(evalcoords) != num_eval) { + PyErr_SetString( + PyExc_ValueError, + "eval_at_coord_proj: eval_dirs and eval_coords must have the same " + "length"); + goto fail; + } + if (ndim_tar < 1 || ndim_tar > gpython_MAX_DIM || + PyArray_SIZE(ncavg) != ndim_tar) { + PyErr_SetString( + PyExc_ValueError, + "eval_at_coord_proj: cells_tar must have between 1 and 7 entries " + "matching ndim_tar"); + goto fail; + } + int out_btype, out_poly_order, out_cdim, out_vdim; + gpython_array *out = gpython_eval_at_coord_proj( + b, cdim_do, ndim, PyArray_DATA(lo), PyArray_DATA(up), PyArray_DATA(nc), + num_eval, PyArray_DATA(evaldirs), PyArray_DATA(evalcoords), ndim_tar, + PyArray_DATA(ncavg), in, &out_btype, &out_poly_order, &out_cdim, + &out_vdim); + Py_DECREF(lo); + Py_DECREF(up); + Py_DECREF(nc); + Py_DECREF(evaldirs); + Py_DECREF(evalcoords); + Py_DECREF(ncavg); + if (!out) { + PyErr_SetString(PyExc_ValueError, + "eval_at_coord_proj: operand shapes incompatible with the " + "basis/grid, or eval_dirs out of range"); + return NULL; + } + PyObject *outcap = wrap_array(out); + if (!outcap) + return NULL; + return Py_BuildValue("(Niiii)", outcap, out_btype, out_poly_order, out_cdim, + out_vdim); +fail: + Py_XDECREF(lo); + Py_XDECREF(up); + Py_XDECREF(nc); + Py_XDECREF(evaldirs); + Py_XDECREF(evalcoords); + Py_XDECREF(ncavg); + return NULL; +} + +/* --------------------------------------------------------- dynvectors */ +static PyObject * +py_dynvec_read(PyObject *self, PyObject *args) +{ + const char *fname; + if (!PyArg_ParseTuple(args, "s", &fname)) + return NULL; + size_t ncomp; + gpython_array *tm = NULL, *data = NULL; + int status = gpython_dynvec_read(fname, &ncomp, &tm, &data); + if (status != 0) { + static const char *msgs[] = { + "", + "no such dynvector file (or empty/unrecognized header)", + "dynvector is not double-precision (unsupported)", + "failed to read dynvector data", + }; + PyErr_Format(PyExc_OSError, "'%s': %s", fname, + msgs[status >= 1 && status <= 3 ? status : 0]); + return NULL; + } + PyObject *tm_cap = wrap_array(tm); + PyObject *data_cap = wrap_array(data); + if (!tm_cap || !data_cap) { + Py_XDECREF(tm_cap); + Py_XDECREF(data_cap); + return NULL; + } + return Py_BuildValue("(nNN)", (Py_ssize_t)ncomp, tm_cap, data_cap); +} + +static PyObject * +py_dynvec_write(PyObject *self, PyObject *args) +{ + const char *fname; + PyObject *tmobj, *dataobj; + if (!PyArg_ParseTuple(args, "sOO", &fname, &tmobj, &dataobj)) + return NULL; + PyArrayObject *tm = + (PyArrayObject *)PyArray_FROM_OTF(tmobj, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY); + PyArrayObject *data = (PyArrayObject *)PyArray_FROM_OTF(dataobj, NPY_DOUBLE, + NPY_ARRAY_IN_ARRAY); + if (!tm || !data) { + Py_XDECREF(tm); + Py_XDECREF(data); + return NULL; + } + npy_intp n = PyArray_DIM(tm, 0); + npy_intp ncomp = PyArray_NDIM(data) > 1 ? PyArray_DIM(data, 1) : 1; + if (PyArray_DIM(data, 0) != n) { + Py_DECREF(tm); + Py_DECREF(data); + PyErr_SetString(PyExc_ValueError, + "dynvec_write: tm and data must share the same length"); + return NULL; + } + /* gpython_dynvec_write (gkeyll/core/zero/dynvec.c) returns raw errno from its + * fopen/fwrite calls, which is only set on failure and never cleared on + * success -- so a stale errno from an earlier, unrelated failed syscall in + * this process would otherwise be misread as this write having failed. */ + errno = 0; + int status = gpython_dynvec_write(fname, (size_t)ncomp, (size_t)n, + PyArray_DATA(tm), PyArray_DATA(data)); + Py_DECREF(tm); + Py_DECREF(data); + if (status != 0) { + PyErr_Format(PyExc_OSError, "'%s': dynvector write failed", fname); + return NULL; + } + Py_RETURN_NONE; +} + +/* --------------------------------------------------------------- module */ +static PyMethodDef gpython_methods[] = { + {"api_version", py_api_version, METH_NOARGS, "gpython shim API version"}, + {"array_new", py_array_new, METH_VARARGS, "zeroed native array"}, + {"array_from_numpy", py_array_from_numpy, METH_VARARGS, + "zero-copy native view of a (…, ncomp) float64 array"}, + {"array_clone", py_array_clone, METH_VARARGS, "deep copy"}, + {"array_ncomp", py_array_ncomp, METH_VARARGS, "components per cell"}, + {"array_size", py_array_size, METH_VARARGS, "number of cells"}, + {"array_view", py_array_view, METH_VARARGS, + "read-only (size, ncomp) view pinning the native memory"}, + {"file_type", py_file_type, METH_VARARGS, "gkyl file type"}, + {"read_header", py_read_header, METH_VARARGS, + "((ndim, lower, upper, cells), file_type, meta, esznc, tot_cells)"}, + {"read_field", py_read_field, METH_VARARGS, + "((ndim, lower, upper, cells), array)"}, + {"basis_new", py_basis_new, METH_VARARGS, "basis handle"}, + {"basis_new_hybrid", py_basis_new_hybrid, METH_VARARGS, + "basis handle (hybrid/gkhybrid, by cdim/vdim)"}, + {"basis_info", py_basis_info, METH_VARARGS, + "(ndim, poly_order, num_basis, id)"}, + {"basis_eval", py_basis_eval, METH_VARARGS, "basis functions at a point"}, + {"basis_node_list", py_basis_node_list, METH_VARARGS, + "(num_basis, ndim) node coordinates"}, + {"basis_nodal_to_modal", py_basis_nodal_to_modal, METH_VARARGS, + "one-cell nodal -> modal"}, + {"dg_mul", py_dg_mul, METH_VARARGS, "weak product (per field)"}, + {"dg_div", py_dg_div, METH_VARARGS, "weak quotient (per field)"}, + {"dg_inv", py_dg_inv, METH_VARARGS, "weak reciprocal (per field)"}, + {"dg_mul_conf_phase", py_dg_mul_conf_phase, METH_VARARGS, + "conf-space x phase-space weak product (single field)"}, + {"array_set", py_array_set, METH_VARARGS, "out = c*a"}, + {"array_accumulate", py_array_accumulate, METH_VARARGS, "out += c*a"}, + {"array_scale", py_array_scale, METH_VARARGS, "a *= c (in place)"}, + {"array_shiftc", py_array_shiftc, METH_VARARGS, + "a[:, comp] += val (in place)"}, + {"array_reduce", py_array_reduce, METH_VARARGS, + "per-component min/max/sum"}, + {"array_dg_reduce", py_array_dg_reduce, METH_VARARGS, + "field-aware (Gauss-node) min/max/sum of one component"}, + {"array_integrate", py_array_integrate, METH_VARARGS, + "int dx op(f) per field"}, + {"array_average", py_array_average, METH_VARARGS, + "weighted (or plain) average over the flagged dims (single field)"}, + {"powsqrt", py_powsqrt, METH_VARARGS, + "pow(sqrt(a), exponent) via gkyl_proj_powsqrt_on_basis (single field)"}, + {"dg_differentiate", py_dg_differentiate, METH_VARARGS, + "local DG derivative (per field, no inter-cell stencil)"}, + {"eval_at_coord_proj", py_eval_at_coord_proj, METH_VARARGS, + "evaluate at coords and project onto the lower-dim target basis; " + "returns (array, target_btype, target_poly_order, target_cdim, " + "target_vdim)"}, + {"write_field", py_write_field, METH_VARARGS, + "write (lower, upper, cells, meta_bytes_or_None, array) to a .gkyl file"}, + {"dynvec_read", py_dynvec_read, METH_VARARGS, + "(ncomp, tm_array, data_array) read from a dynvector file"}, + {"dynvec_write", py_dynvec_write, METH_VARARGS, + "write a dynvector from parallel tm[n]/data[n,ncomp] arrays"}, + {NULL, NULL, 0, NULL}, +}; + +static struct PyModuleDef gpython_module = { + PyModuleDef_HEAD_INIT, + "_gpython", + "Compiled bridge to Gkeyll via the gpython shim (see GKEYLL_C_SHIM.md).", + -1, + gpython_methods, +}; + +PyMODINIT_FUNC +PyInit__gpython(void) +{ + import_array(); + PyObject *m = PyModule_Create(&gpython_module); + if (!m) + return NULL; + if (PyModule_AddIntConstant(m, "GPYTHON_API_VERSION", GPYTHON_API_VERSION) < + 0) { + Py_DECREF(m); + return NULL; + } + return m; +} diff --git a/src/postgkyl/gpython/kernels.py b/src/postgkyl/gpython/kernels.py new file mode 100644 index 00000000..4bb9e06c --- /dev/null +++ b/src/postgkyl/gpython/kernels.py @@ -0,0 +1,641 @@ +"""Thin wrappers over Gkeyll's compiled operators (weak algebra & reductions). + +Each function takes :class:`~postgkyl.gpython.array.GkylArray` operands plus the +basis descriptor and calls one shim entry point; the per-field loop for +``ncomp == nfields * num_basis`` arrays and all transient C resources +(``gkyl_dg_bin_op_mem``, integrate updaters) live inside the compiled shim. + +Python-side capability guards mirror Gkeyll's own limits (which are C +``assert``s -- letting them fire would abort the process). +""" + +from __future__ import annotations + +import numpy as np + +from . import _lib +from .array import GkylArray +from .basis import get_basis + +_WEAK_BASES = ("serendipity", "tensor") # dg_bin_ops: assert(false) otherwise + +# enum gkyl_array_op / gkyl_array_integrate_op ordinals used by the shim +REDUCE_OPS = {"min": 0, "max": 1, "sum": 2} +GKYL_MIN, GKYL_MAX, GKYL_SUM = 0, 1, 2 +INTEGRATE_OPS = {"none": 0, "abs": 1, "sq": 2} + +# Weak mul/div kernel tables (gkyl_dg_bin_ops_priv.h ser_mul_list/ten_mul_list/ +# ser_div_set_list/ten_div_set_list) are fixed-size [ndim][poly_order] arrays +# covering ONLY ndim 1..3 -- narrower than the basis module's own eval range. +# ndim >= 4 hits `assert(dim < 4)` in choose_ser_mul_kern (a process abort); +# an out-of-table poly_order for tensor (p3 at ndim 2-3) returns a NULL +# kernel pointer that gkyl_dg_mul_op/div_op call with NO null check at all +# (a segfault, not an assert). Both must be refused here. +_WEAK_MAX_POLY_ORDER = { + "serendipity": { + 1: 3, + 2: 3, + 3: 3 + }, + "tensor": { + 1: 3, + 2: 2, + 3: 2 + }, +} +# gkyl_dg_inv_op's kernel table (ser_inv_list) has no dim bound check +# whatsoever (a raw out-of-bounds array read for ndim >= 4) and only fills +# poly_order == 1 for ndim 1..3. +_WEAK_INV_DIMS = (1, 2, 3) + + +def _check_weak(basis_type: str, ndim: int, poly_order: int, *arrays: + GkylArray): + basis_type = basis_type.lower() + limits = _WEAK_MAX_POLY_ORDER.get(basis_type) + if limits is None: + raise NotImplementedError( + f"Gkeyll weak ops support {_WEAK_BASES}, not '{basis_type}'") + max_p = limits.get(ndim) + if max_p is None: + raise NotImplementedError( + f"Gkeyll's weak (DG) mul/div kernels support ndim 1..3, got {ndim}") + if not 0 <= poly_order <= max_p: + raise NotImplementedError( + f"Gkeyll's weak {basis_type} mul/div kernels in {ndim}D support " + f"poly_order 0..{max_p}, got {poly_order}") + first = arrays[0] + for a in arrays[1:]: + if (a.ncomp, a.size) != (first.ncomp, first.size): + raise ValueError(f"operand shape mismatch: {a.ncomp}x{a.size} vs " + f"{first.ncomp}x{first.size}") + + +def _fields(arr: GkylArray, num_basis: int) -> int: + if arr.ncomp % num_basis: + raise ValueError(f"ncomp {arr.ncomp} is not a multiple of " + f"num_basis {num_basis}") + return arr.ncomp // num_basis + + +def weak_mul(basis_type: str, ndim: int, poly_order: int, a: GkylArray, + b: GkylArray) -> GkylArray: + """Weak (DG) product ``a * b``, field by field, via ``gkyl_dg_mul_op``.""" + _check_weak(basis_type, ndim, poly_order, a, b) + basis = get_basis(basis_type, ndim, poly_order) + _fields(a, basis.num_basis) + out = GkylArray.alloc(a.ncomp, a.size) + _lib.require().dg_mul(basis._cap, out._cap, a._cap, b._cap) + return out + + +def weak_div(basis_type: str, ndim: int, poly_order: int, a: GkylArray, + b: GkylArray) -> GkylArray: + """Weak (DG) quotient ``a / b`` via ``gkyl_dg_div_op`` (per-cell solve).""" + _check_weak(basis_type, ndim, poly_order, a, b) + basis = get_basis(basis_type, ndim, poly_order) + _fields(a, basis.num_basis) + out = GkylArray.alloc(a.ncomp, a.size) + _lib.require().dg_div(basis._cap, out._cap, a._cap, b._cap) + return out + + +def weak_inv(basis_type: str, ndim: int, poly_order: int, + a: GkylArray) -> GkylArray: + """Weak reciprocal ``1 / a`` via ``gkyl_dg_inv_op`` (Gkeyll: ser p=1, ndim<=3 only).""" + if basis_type.lower() != "serendipity" or poly_order != 1: + raise NotImplementedError( + "gkyl_dg_inv_op supports serendipity p=1 only (a Gkeyll limit); " + "use weak division instead.") + if ndim not in _WEAK_INV_DIMS: + raise NotImplementedError( + f"gkyl_dg_inv_op supports ndim {_WEAK_INV_DIMS} only, got {ndim} " + "(a Gkeyll limit; its kernel table has no bounds check at all, so " + "this guard is load-bearing, not decorative)") + basis = get_basis(basis_type, ndim, poly_order) + _fields(a, basis.num_basis) + out = GkylArray.alloc(a.ncomp, a.size) + _lib.require().dg_inv(basis._cap, out._cap, a._cap) + return out + + +# -------------------------------------------------- conf-space x phase-space +# gkyl_dg_mul_conf_phase_op_range picks its kernel from the PHASE basis type +# alone (choose_mul_conf_phase_kern in gkyl_dg_bin_ops_priv.h); for +# hybrid/gkhybrid the conf poly_order it reads is unused by that branch, so +# the only real requirement (Gkeyll's own PKPM/GK convention) is a +# serendipity conf basis. Every (cdim, vdim) split our own basis.py +# convention (_HYBRID_CDIM_VDIM) derives from a valid hybrid/gkhybrid ndim +# already has a populated cross_mul_list entry -- verified by hand against +# hyb_cross_mul_list/gkhyb_cross_mul_list, so no extra table is needed there. +# serendipity/tensor phase bases have a genuinely holey (cdim, pdim, +# poly_order) kernel table -- unlike same-basis weak_mul, most combinations +# a valid same-basis object could have are simply absent (NULL function +# pointer, no null check in gkyl_dg_mul_conf_phase_op_range) -- so those are +# guarded explicitly below, transcribed from ser_cross_mul_list / +# ten_cross_mul_list in gkyl_dg_bin_ops_priv.h. +_CROSS_MUL_SER = { + 2: { + 1: {1, 2, 3} + }, + 3: { + 1: {1, 2, 3}, + 2: {1, 2, 3} + }, + 4: { + 1: {1, 2, 3}, + 2: {1, 2, 3}, + 3: {1, 2, 3} + }, + 5: { + 2: {1, 2}, + 3: {1, 2} + }, + 6: { + 3: {1} + }, +} +_CROSS_MUL_TEN = { + 2: { + 1: {1, 2} + }, + 3: { + 1: {1, 2}, + 2: {1, 2} + }, + 4: { + 1: {1, 2}, + 2: {1, 2}, + 3: {1} + }, + 5: { + 2: {1, 2}, + 3: {1} + }, + 6: { + 3: {1} + }, +} +_CROSS_MUL_TABLES = {"serendipity": _CROSS_MUL_SER, "tensor": _CROSS_MUL_TEN} + + +def _check_mul_conf_phase(conf_basis_type: str, phase_basis_type: str, + conf_ndim: int, phase_ndim: int, poly_order: int): + conf_basis_type = conf_basis_type.lower() + phase_basis_type = phase_basis_type.lower() + if phase_ndim <= conf_ndim: + raise ValueError( + f"phase_ndim ({phase_ndim}) must exceed conf_ndim ({conf_ndim})") + if phase_basis_type in ("hybrid", "gkhybrid"): + if conf_basis_type != "serendipity": + raise NotImplementedError( + "Gkeyll pairs a hybrid/gkhybrid phase basis with a serendipity " + f"conf basis only (its own PKPM/GK convention), not " + f"'{conf_basis_type}'") + return + if phase_basis_type in ("serendipity", "tensor"): + if conf_basis_type != phase_basis_type: + raise NotImplementedError( + "gkyl_dg_mul_conf_phase_op_range picks its kernel from the phase " + f"basis type alone ('{phase_basis_type}'); pair it with a conf " + f"basis of the same type, not '{conf_basis_type}'") + valid = _CROSS_MUL_TABLES[phase_basis_type].get(phase_ndim, + {}).get(conf_ndim) + if not valid or poly_order not in valid: + raise NotImplementedError( + f"Gkeyll has no {phase_basis_type} conf*phase cross-mul kernel " + f"for conf_ndim={conf_ndim}, phase_ndim={phase_ndim}, " + f"poly_order={poly_order}") + return + raise NotImplementedError( + "Gkeyll's conf*phase cross-mul supports serendipity, tensor, hybrid, " + f"gkhybrid phase bases, not '{phase_basis_type}'") + + +def weak_mul_conf_phase(conf_basis_type: str, conf_ndim: int, + phase_basis_type: str, phase_ndim: int, poly_order: int, + conf_cells, phase_cells, cop: GkylArray, + pop: GkylArray) -> GkylArray: + """Conf-space x phase-space weak product ``cop * pop`` via + ``gkyl_dg_mul_conf_phase_op_range`` -- e.g. a density (conf-space) times a + distribution function (phase-space) in PKPM/gyrokinetic post-processing. + + Unlike :func:`weak_mul`, this is single-field only on both sides (the + underlying kernel takes no field-index arguments): ``cop.ncomp`` must + equal the conf basis's ``num_basis`` and ``pop.ncomp`` the phase basis's. + + ``conf_cells``/``phase_cells`` are each grid's per-dimension cell counts + (e.g. ``rio``'s ``grid["cells"]``) -- Gkeyll maps each phase cell to its + conf cell by dropping the velocity-space indices, so both cell counts are + needed to build matching index ranges; ``conf_cells`` must equal the + leading ``conf_ndim`` entries of ``phase_cells``. + + The dispatch is asymmetric: Gkeyll chooses the kernel from the PHASE + basis type alone, so ``conf_basis_type`` must be ``"serendipity"`` when + pairing with hybrid/gkhybrid, or match ``phase_basis_type`` exactly for + serendipity/tensor. + """ + _check_mul_conf_phase(conf_basis_type, phase_basis_type, conf_ndim, + phase_ndim, poly_order) + cbasis = get_basis(conf_basis_type, conf_ndim, poly_order) + pbasis = get_basis(phase_basis_type, phase_ndim, poly_order) + if cop.ncomp != cbasis.num_basis: + raise ValueError( + f"cop.ncomp ({cop.ncomp}) must equal the conf basis's num_basis " + f"({cbasis.num_basis}); mul_conf_phase is single-field only") + if pop.ncomp != pbasis.num_basis: + raise ValueError( + f"pop.ncomp ({pop.ncomp}) must equal the phase basis's num_basis " + f"({pbasis.num_basis}); mul_conf_phase is single-field only") + conf_cells = np.asarray(conf_cells, dtype=np.int32) + phase_cells = np.asarray(phase_cells, dtype=np.int32) + out = GkylArray.alloc(pop.ncomp, pop.size) + _lib.require().dg_mul_conf_phase(cbasis._cap, pbasis._cap, out._cap, cop._cap, + pop._cap, conf_cells, phase_cells) + return out + + +# ------------------------------------------------------- linear coefficient ops +def lincomb(ca: float, a: GkylArray, cb: float, b: GkylArray) -> GkylArray: + """``ca*a + cb*b`` on the DG coefficients (gkyl_array_set + accumulate).""" + if (a.ncomp, a.size) != (b.ncomp, b.size): + raise ValueError("operand shape mismatch in lincomb") + g0 = _lib.require() + out = GkylArray.alloc(a.ncomp, a.size) + g0.array_set(out._cap, ca, a._cap) + g0.array_accumulate(out._cap, cb, b._cap) + return out + + +def scale(a: GkylArray, factor: float) -> GkylArray: + """``factor * a`` (gkyl_array_scale on a clone; the input is untouched).""" + out = a.clone() + _lib.require().array_scale(out._cap, factor) + return out + + +def shiftc(a: GkylArray, val: float, comp: int) -> GkylArray: + """Add ``val`` to component ``comp`` of every cell (gkyl_array_shiftc).""" + out = a.clone() + _lib.require().array_shiftc(out._cap, float(val), comp) + return out + + +# ---------------------------------------------------------------- reductions +def reduce(a: GkylArray, op: int) -> np.ndarray: + """Per-component MIN/MAX/SUM over all cells (gkyl_array_reduce). + + This reduces the raw DG **coefficients**: exact for ``"sum"`` (the sum of + coefficients over cells is linear), but NOT the field's true min/max -- a + DG expansion can exceed its coefficient values between nodes. Use + :func:`dg_reduce` for the field-aware version. + """ + return _lib.require().array_reduce(a._cap, op) + + +def dg_reduce(basis_type: str, ndim: int, poly_order: int, a: GkylArray, + comp: int, op: str) -> float: + """MIN/MAX/SUM of the field ``comp`` actually represents (gkyl_array_dg_reducec). + + Evaluates the DG expansion at each cell's Gauss-Legendre quadrature nodes + and reduces those values -- the true min/max/sum of the represented field, + exact for ``"sum"`` and correct (not merely coefficient-bounded) for + ``"min"``/``"max"`` to quadrature precision (exact for polynomials the + quadrature integrates exactly, i.e. always for a basis's own degree). + + Args: + basis_type: ``"serendipity"`` or ``"tensor"``. + ndim: number of dimensions the basis was built for. + poly_order: polynomial order the basis was built for. + a: array whose ``ncomp`` is a multiple of the basis's ``num_basis``. + comp: 0-based field index (NOT a coefficient offset). + op: one of ``"min"``, ``"max"``, ``"sum"``. + + Returns: + The reduced scalar. + + Raises: + ValueError: unknown ``op``, or ``comp`` out of range for ``a``'s fields. + """ + if op not in REDUCE_OPS: + raise ValueError(f"dg_reduce op '{op}' not in {sorted(REDUCE_OPS)}") + basis = get_basis(basis_type, ndim, poly_order) + nfields = _fields(a, basis.num_basis) + if not 0 <= comp < nfields: + raise ValueError(f"comp {comp} out of range for {nfields} field(s)") + return float(_lib.require().array_dg_reduce(basis._cap, a._cap, comp, + REDUCE_OPS[op])) + + +def array_average(grid: dict, + basis_type: str, + poly_order: int, + ndim_avg: int, + cells_avg, + avg_dim, + a: GkylArray, + weight: GkylArray | None = None) -> GkylArray: + """Single-field weighted (or plain) average of ``a`` via ``gkyl_array_average``: + ``int f w dx^avg / int w dx^avg`` (or ``int f dx^avg / int dx^avg`` when + ``weight`` is omitted). + + ``grid`` is the donor grid dict (ndim/lower/upper/cells, e.g. from ``rio``); + ``avg_dim`` flags (length ``grid["ndim"]``, 1 = averaged, 0 = kept) which + donor dims are reduced. ``ndim_avg``/``cells_avg`` describe the target: the + surviving dims' cell counts, or ``ndim_avg=1``/``cells_avg=[1]`` for a full + reduction (Gkeyll's own convention -- there is no true 0-dimensional + basis). Single-field only: ``a``/``weight`` must carry exactly one donor + basis's worth of coefficients (``a.ncomp == num_basis``); a multi-field + caller loops field by field (:func:`postgkyl.dg.modal.average`). + + Guarded to the kernel set compiled into libg0core (serendipity p1-p2, + donor ndim 1-3 -- ``gkyl_array_average_new`` asserts ``poly_order <= 2`` + and its kernel-choice table has no bound check past that). + """ + if basis_type.lower() != "serendipity" or poly_order not in (1, 2): + raise NotImplementedError( + "gkyl_array_average kernels in libg0core cover serendipity p1-p2") + ndim = int(grid["ndim"]) + if ndim not in (1, 2, 3): + raise NotImplementedError( + f"gkyl_array_average kernels in libg0core cover ndim 1-3, got {ndim}") + basis = get_basis(basis_type, ndim, poly_order) + basis_avg = get_basis(basis_type, ndim_avg, poly_order) + if a.ncomp != basis.num_basis: + raise ValueError( + f"average: a.ncomp ({a.ncomp}) must equal the donor basis's " + f"num_basis ({basis.num_basis}); average is single-field only") + if weight is not None and (weight.ncomp, weight.size) != (basis.num_basis, + a.size): + raise ValueError( + f"average: weight must be single-field ({basis.num_basis} comps) " + f"and share the donor array's size ({a.size} cells)") + lower = np.asarray(grid["lower"], dtype=np.float64) + upper = np.asarray(grid["upper"], dtype=np.float64) + cells = np.asarray(grid["cells"], dtype=np.int32) + if int(np.prod(cells)) != a.size: + raise ValueError(f"grid cells {tuple(cells)} do not cover the array " + f"({int(np.prod(cells))} vs {a.size} cells)") + cells_avg = np.asarray(cells_avg, dtype=np.int32) + avg_dim = np.asarray(avg_dim, dtype=np.int32) + out = GkylArray.alloc(basis_avg.num_basis, int(np.prod(cells_avg))) + _lib.require().array_average(lower, upper, cells, basis._cap, basis_avg._cap, + cells_avg, avg_dim, + weight._cap if weight is not None else None, + a._cap, out._cap) + return out + + +def integrate(grid: dict, + basis_type: str, + poly_order: int, + a: GkylArray, + op: str = "none", + factor: float = 1.0) -> np.ndarray: + """``int dx op(f)`` per field via ``gkyl_array_integrate`` -- one value per field. + + ``grid`` is the dict from ``rio`` (ndim/lower/upper/cells). Guarded to the + kernel set compiled into libg0core (serendipity p1-p2, ndim 1-3, for + none/abs/sq) -- ``gkyl_array_integrate_choose_kernel`` indexes its kernel + table by ``ndim-1``/``poly_order-1`` with no bound past an + ``assert(up->kernel)`` that a genuinely out-of-table ndim can dodge (an + out-of-bounds array read that happens to be non-NULL), so ndim is checked + here rather than left to that assert. + """ + if op not in INTEGRATE_OPS: + raise ValueError(f"integrate op '{op}' not in {sorted(INTEGRATE_OPS)}") + if basis_type.lower() != "serendipity" or poly_order not in (1, 2): + raise NotImplementedError( + "gkyl_array_integrate kernels in libg0core cover serendipity p1-p2") + ndim = int(grid["ndim"]) + if ndim not in (1, 2, 3): + raise NotImplementedError( + f"gkyl_array_integrate kernels in libg0core cover ndim 1-3, got {ndim}") + basis = get_basis(basis_type, ndim, poly_order) + nfields = _fields(a, basis.num_basis) + lower = np.asarray(grid["lower"], dtype=np.float64) + upper = np.asarray(grid["upper"], dtype=np.float64) + cells = np.asarray(grid["cells"], dtype=np.int32) + if int(np.prod(cells)) != a.size: + raise ValueError(f"grid cells {tuple(cells)} do not cover the array " + f"({int(np.prod(cells))} vs {a.size} cells)") + return _lib.require().array_integrate(lower, upper, cells, + basis._cap, nfields, INTEGRATE_OPS[op], + float(factor), a._cap) + + +# -------------------------------------------------------------------- powsqrt +def powsqrt(basis_type: str, + ndim: int, + poly_order: int, + cells, + a: GkylArray, + exponent: float, + num_quad: int | None = None) -> GkylArray: + """Single-field ``pow(sqrt(a), exponent)`` (i.e. ``a ** (exponent/2)``) via + ``gkyl_proj_powsqrt_on_basis`` -- a Gauss-Legendre-quadrature projection, + not a fixed per-(basis_type, ndim, poly_order) kernel table like every + other function in this module: the real updater works off the basis's own + ``eval`` callback, so there is no coverage guard here beyond the shape + check below. + + A negative value at a quadrature node is clamped to ``1e-40`` by the + updater itself (Gkeyll's own convention), which can differ from the DG + coefficients' own sign between nodes -- this is not re-validated here. + + ``cells`` is the grid's per-dimension cell count (e.g. ``ctx["cells"]``); + no physical extent is needed, only cell indexing. ``num_quad`` defaults to + ``poly_order + 1``, matching the gyrokinetic app's own use of this + updater. + """ + basis = get_basis(basis_type, ndim, poly_order) + if a.ncomp != basis.num_basis: + raise ValueError( + f"powsqrt: a.ncomp ({a.ncomp}) must equal the basis's num_basis " + f"({basis.num_basis}); powsqrt is single-field only") + if num_quad is None: + num_quad = poly_order + 1 + if num_quad < 1: + raise ValueError(f"num_quad must be >= 1, got {num_quad}") + cells = np.asarray(cells, dtype=np.int32) + if int(np.prod(cells)) != a.size: + raise ValueError(f"cells {tuple(cells)} do not cover the array " + f"({int(np.prod(cells))} vs {a.size} cells)") + out = GkylArray.alloc(basis.num_basis, a.size) + _lib.require().powsqrt(basis._cap, int(num_quad), float(exponent), cells, + out._cap, a._cap) + return out + + +# --------------------------------------------------------------- differentiate +# gkyl_dg_differentiate_op_local's kernel tables (gkyl_dg_differentiate_priv.h +# ser_differentiate_list/ten_differentiate_list) cover only serendipity and +# tensor, with NO bounds check at all in the dispatch (an unconditional +# `assert(diff_op)` on a NULL table entry -- a process abort, not a Python +# exception), so this guard must run before every call. +_DIFFERENTIATE_MAX_POLY_ORDER = { + "serendipity": { + 1: 2, + 2: 2, + 3: 1 + }, + "tensor": { + 1: 2, + 2: 2 + }, # ndim 3: no tensor differentiate kernels at all +} + + +def _check_differentiate(basis_type: str, ndim: int, poly_order: int, dir: int, + diff_order: int): + basis_type = basis_type.lower() + limits = _DIFFERENTIATE_MAX_POLY_ORDER.get(basis_type) + if limits is None: + raise NotImplementedError( + f"gkyl_dg_differentiate_op_local supports serendipity/tensor, not " + f"'{basis_type}'") + max_p = limits.get(ndim) + if max_p is None: + raise NotImplementedError( + f"Gkeyll's {basis_type} differentiate kernels support ndim " + f"{sorted(limits)}, got {ndim}") + if not 1 <= poly_order <= max_p: + raise NotImplementedError( + f"Gkeyll's {basis_type} differentiate kernels in {ndim}D support " + f"poly_order 1..{max_p}, got {poly_order}") + if not 0 <= dir < ndim: + raise ValueError(f"differentiate dir {dir} out of range for a {ndim}D " + "field") + if diff_order not in (1, 2): + raise ValueError(f"differentiate order must be 1 or 2, got {diff_order}") + + +def weak_differentiate(basis_type: str, ndim: int, poly_order: int, dir: int, + diff_order: int, dx: float, a: GkylArray) -> GkylArray: + """Local DG derivative ``d^diff_order/dx_dir^diff_order a`` via + ``gkyl_dg_differentiate_op_local``, field by field. + + Differentiates the DG expansion independently in every cell (no + inter-cell stencil) -- an exact derivative of the polynomial each cell + already represents, not a finite-difference approximation across cells. + Serendipity/tensor only (a Gkeyll limit); ``dx`` is the cell length along + ``dir``. + """ + _check_differentiate(basis_type, ndim, poly_order, dir, diff_order) + basis = get_basis(basis_type, ndim, poly_order) + out = GkylArray.alloc(a.ncomp, a.size) + _lib.require().dg_differentiate(basis._cap, dir, diff_order, float(dx), + out._cap, a._cap) + return out + + +# ------------------------------------------------------- evaluate-and-project +# gkyl_dg_eval_at_coord_proj's own dispatch (gkyl_dg_eval_at_coord_proj_priv.h) +# covers serendipity (ndim 1-4 at p1-p2, ndim 5-6 at p1 only), tensor (ndim +# 1 and 3 at p1 only, ndim 2 at p1-p2), and gkhybrid (p1 only, at the same +# (cdim, vdim) combinations basis.py's own hybrid table already recognizes). +# Plain "hybrid" is not in that dispatch's switch at all. Like the tables +# above, an out-of-coverage combination is a process abort (an unconditional +# `assert(kers->ev_ker)` on a NULL table entry), not a Python exception. +_EVAL_AT_COORD_PROJ_MAX_POLY_ORDER = { + "serendipity": { + 1: 2, + 2: 2, + 3: 2, + 4: 2, + 5: 1, + 6: 1 + }, + "tensor": { + 1: 1, + 2: 2, + 3: 1 + }, +} + +# gkyl_basis_type ordinals (gkeyll/core/zero/gkyl_basis.h) -- the target +# basis gpython_eval_at_coord_proj reports can differ in TYPE from the donor +# (e.g. eliminating a gkhybrid velocity direction can yield a plain +# serendipity target), so its ordinal must be translated back to postgkyl's +# string vocabulary here. +_BASIS_TYPE_ORDINALS = { + 0: "serendipity", + 1: "tensor", + 2: "hybrid", + 3: "gkhybrid", + 4: "gkhybrid_vel", +} + + +def _check_eval_at_coord_proj(basis_type: str, ndim: int, poly_order: int, + eval_dirs): + basis_type = basis_type.lower() + if basis_type == "gkhybrid": + if poly_order != 1: + raise NotImplementedError( + "gkyl_dg_eval_at_coord_proj's gkhybrid kernels exist at " + f"poly_order 1 only, got {poly_order}") + else: + limits = _EVAL_AT_COORD_PROJ_MAX_POLY_ORDER.get(basis_type) + if limits is None: + raise NotImplementedError( + "gkyl_dg_eval_at_coord_proj supports serendipity/tensor/gkhybrid, " + f"not '{basis_type}'") + max_p = limits.get(ndim) + if max_p is None: + raise NotImplementedError( + f"Gkeyll's {basis_type} eval_at_coord_proj kernels support ndim " + f"{sorted(limits)}, got {ndim}") + if not 1 <= poly_order <= max_p: + raise NotImplementedError( + f"Gkeyll's {basis_type} eval_at_coord_proj kernels in {ndim}D " + f"support poly_order 1..{max_p}, got {poly_order}") + eval_dirs = sorted(set(int(d) for d in eval_dirs)) + if not eval_dirs or eval_dirs[0] < 0 or eval_dirs[-1] >= ndim: + raise ValueError(f"eval_dirs {eval_dirs} out of range for a {ndim}D " + "field") + return eval_dirs + + +def eval_at_coord_proj(basis_type: str, ndim: int, poly_order: int, + cdim_do: int, grid: dict, eval_dirs, eval_coords, + ndim_tar: int, cells_tar, a: GkylArray): + """Evaluate ``a`` at ``eval_coords`` in ``eval_dirs`` and project onto the + lower-dimensional target basis Gkeyll picks for that elimination, via + ``gkyl_dg_eval_at_coord_proj``. + + ``grid`` is the donor grid dict (``ndim``/``lower``/``upper``/``cells``, + e.g. from ``rio``). ``cdim_do`` is the donor's configuration-space + dimension count (equal to ``ndim`` for serendipity/tensor; the (cdim, + vdim) split for gkhybrid -- see ``basis._HYBRID_CDIM_VDIM``). + ``ndim_tar``/``cells_tar`` describe the target's rectangular index range + with the same convention :func:`array_average` uses: the surviving donor + dims' cell counts in donor order, or ``ndim_tar=1``/``cells_tar=[1]`` for + a full reduction (every donor direction evaluated away). + + Returns: + ``(out, target_basis_type, target_poly_order, target_cdim, + target_vdim)`` -- the target array, field by field like the donor + (``ncomp`` scaled to the donor's field count), and the target basis's + metadata (which can differ in TYPE from the donor's, e.g. eliminating a + gkhybrid velocity direction can yield a plain serendipity target). + """ + eval_dirs = _check_eval_at_coord_proj(basis_type, ndim, poly_order, eval_dirs) + basis = get_basis(basis_type, ndim, poly_order) + lower = np.asarray(grid["lower"], dtype=np.float64) + upper = np.asarray(grid["upper"], dtype=np.float64) + cells = np.asarray(grid["cells"], dtype=np.int32) + if int(np.prod(cells)) != a.size: + raise ValueError(f"grid cells {tuple(cells)} do not cover the array " + f"({int(np.prod(cells))} vs {a.size} cells)") + eval_dirs_arr = np.asarray(eval_dirs, dtype=np.int32) + eval_coords_arr = np.asarray(eval_coords, dtype=np.float64) + if eval_coords_arr.shape != eval_dirs_arr.shape: + raise ValueError("eval_dirs and eval_coords must have the same length") + cells_tar = np.asarray(cells_tar, dtype=np.int32) + out_cap, btype, poly_order_tar, cdim_tar, vdim_tar = ( + _lib.require().eval_at_coord_proj(basis._cap, int(cdim_do), lower, upper, + cells, eval_dirs_arr, eval_coords_arr, + int(ndim_tar), cells_tar, a._cap)) + return (GkylArray(out_cap), _BASIS_TYPE_ORDINALS[btype], poly_order_tar, + cdim_tar, vdim_tar) diff --git a/src/postgkyl/gpython/rio.py b/src/postgkyl/gpython/rio.py new file mode 100644 index 00000000..cb151e92 --- /dev/null +++ b/src/postgkyl/gpython/rio.py @@ -0,0 +1,123 @@ +"""File loading through Gkeyll's ``gkyl_array_rio`` -- the C read path. + +``read_field`` performs the whole read (grid + allocate + fill, including +multi-range stitching for file_type 3) inside Gkeyll; ``read_header`` returns +the grid and the raw msgpack metadata blob without touching the payload. +Decoding the msgpack bytes is left to the caller (``io/``) -- metadata policy +is an io concern, bytes are a floor concern. +""" + +from __future__ import annotations + +import numpy as np + +from . import _lib +from .array import GkylArray + +# enum gkyl_file_type ordinals used by gkyl_get_gkyl_file_type +FIELD_FILE_TYPES = (1, 3) # single-range and multi-range field data + + +def file_type(file_name: str) -> int: + """The gkyl file type (1..5), or -1 if not a gkyl file.""" + return int(_lib.require().file_type(file_name)) + + +def read_header(file_name: str): + """Header-only read: ``(grid_dict, file_type, meta_bytes, esznc, tot_cells)``. + + ``grid_dict`` has ``ndim``/``lower``/``upper``/``cells`` as NumPy values; + ``meta_bytes`` is the raw msgpack blob (b"" when the file has none). + """ + grid, ftype, meta, esznc, tot_cells = _lib.require().read_header(file_name) + return _grid_dict(grid), ftype, meta, esznc, tot_cells + + +def read_field(file_name: str): + """Full field read inside Gkeyll: ``(grid_dict, GkylArray)``.""" + grid, cap = _lib.require().read_field(file_name) + return _grid_dict(grid), GkylArray(cap) + + +def write_field(file_name: str, + grid: dict, + arr: GkylArray, + *, + meta: bytes = b"") -> None: + """Write ``arr`` on a uniform ``grid`` through ``gkyl_grid_sub_array_write``. + + The same C write path Gkeyll itself uses, so a round trip through this + function and :func:`read_field` is bit-exact by construction. ``meta`` is + a raw msgpack byte blob (encoding policy belongs to ``io/``, which decodes + it the same way on read); pass ``b""`` for no metadata. + + Args: + file_name: destination path. + grid: a dict with ``lower``/``upper``/``cells`` (as returned by + :func:`read_header`/:func:`read_field`, or built by the caller). + arr: the array to write; ``arr.size`` must equal ``prod(grid["cells"])``. + meta: raw msgpack bytes, or empty for none. + + Raises: + ValueError: ``grid["cells"]`` does not cover ``arr``. + OSError: the underlying ``gkyl_array_rio`` write failed. + """ + lower = np.asarray(grid["lower"], dtype=np.float64) + upper = np.asarray(grid["upper"], dtype=np.float64) + cells = np.asarray(grid["cells"], dtype=np.int32) + if int(np.prod(cells)) != arr.size: + raise ValueError(f"grid cells {tuple(cells)} do not cover the array " + f"({int(np.prod(cells))} vs {arr.size} cells)") + _lib.require().write_field(file_name, lower, upper, cells, + meta if meta else None, arr._cap) + + +def read_dynvec(file_name: str): + """Read a time-series (dynvector) file: ``(time (n,), data (n, ncomp))``. + + Args: + file_name: path to a gkyl dynvector file (``file_type`` 2). + + Returns: + ``time``: 1-D array of ``n`` timestamps. + ``data``: ``(n, ncomp)`` array of the recorded values. + + Raises: + OSError: missing file, non-double dynvector, or a read failure. + """ + ncomp, tm_cap, data_cap = _lib.require().dynvec_read(file_name) + tm = GkylArray(tm_cap).to_numpy()[:, 0] + data = GkylArray(data_cap).to_numpy() + return tm, data + + +def write_dynvec(file_name: str, time: np.ndarray, data: np.ndarray) -> None: + """Write a time-series (dynvector) file via ``gkyl_dynvec_write``. + + Args: + file_name: destination path. + time: 1-D array of ``n`` timestamps. + data: ``(n,)`` or ``(n, ncomp)`` array of values, one row per timestamp. + + Raises: + ValueError: ``time`` and ``data`` disagree on the number of samples. + OSError: the underlying ``gkyl_dynvec_write`` failed. + """ + time = np.ascontiguousarray(time, dtype=np.float64) + data = np.asarray(data, dtype=np.float64) + if data.ndim == 1: + data = data[:, None] + if data.shape[0] != time.shape[0]: + raise ValueError(f"time has {time.shape[0]} samples but data has " + f"{data.shape[0]}") + _lib.require().dynvec_write(file_name, time, np.ascontiguousarray(data)) + + +def _grid_dict(grid: tuple) -> dict: + ndim, lower, upper, cells = grid + return { + "ndim": int(ndim), + "lower": np.asarray(lower), + "upper": np.asarray(upper), + "cells": np.asarray(cells, dtype=np.int64), + } diff --git a/src/postgkyl/io/__init__.py b/src/postgkyl/io/__init__.py new file mode 100644 index 00000000..fea72bf9 --- /dev/null +++ b/src/postgkyl/io/__init__.py @@ -0,0 +1,69 @@ +"""File I/O -- bytes <-> dataset arrays. + +A leaf layer: one reader per format, dispatched by ``read()``; ``write()`` for +output. Nothing here imports ``gdatastate``/``operations``; the readers fill a plain ``ctx`` +dict and return ``(grid, values)`` so the container can construct itself on top. +""" + +from __future__ import annotations + +import os.path + +from . import mapping +from .naming import OutputName, parse_output_name +from .gkyl_c_reader import GkylCReader +from .gkyl_reader import GkylReader +from .gkyl_h5_reader import GkylH5Reader +from .flash_h5_reader import FlashH5Reader +from .writer import save + +# Reader registry -- tried in order; extend by adding (name, reader) entries. +# Order is by *specificity* of ``is_compatible()``, most specific / cheapest +# first, so a file never falls into the wrong reader: +# 1. "gkyl_c" -- native .gkyl via libg0core; the magic-byte + file-type +# check is exact and returns modal data as a GkylArray. +# 2. "gkyl" -- pure-Python .gkyl fallback (no libg0core, partial loads, +# dynvectors); same exact magic-byte check as gkyl_c. +# 3. "h5" -- legacy Gkeyll HDF5 output (predates the native .gkyl +# binary format); is_compatible() requires the Gkeyll-specific +# "/StructGridField" or "/DataStruct/data" node, so a FLASH +# .h5 file (no such nodes) is correctly declined and falls +# through to "flash". +# 4. "flash" -- FLASH code HDF5 output; is_compatible() requires a +# "coordinates" node, disjoint from the Gkeyll h5 layout. +# Because 1-2 are checked with the same fast magic-byte test before 3-4 ever +# touch the (slower) tables importer, a .gkyl file never reaches an h5 +# reader, and vice versa. +_READERS = { + "gkyl_c": GkylCReader, + "gkyl": GkylReader, + "h5": GkylH5Reader, + "flash": FlashH5Reader, +} + + +def read(file_name: str, ctx: dict | None = None, **kwargs): + """Read ``file_name`` into ``(grid, values)``, populating ``ctx`` in place. + + The reader is chosen by trying each registered reader's ``is_compatible`` + check. ``ctx`` (a plain dict) is filled with metadata -- ``poly_order``, + ``basis_type``, ``cells``, ``lower``/``upper``, ``time``/``frame``, ... -- + exactly as the legacy reader did. + """ + if ctx is None: + ctx = {} + if not os.path.exists(file_name): + raise FileNotFoundError(f"No such file: '{file_name}'") + for reader_cls in _READERS.values(): + reader = reader_cls(file_name=file_name, ctx=ctx, **kwargs) + if reader.is_compatible(): + reader.preload() + return reader.load() + raise NameError( + f"'{file_name}' cannot be read with any known reader: {list(_READERS)}") + + +__all__ = [ + "read", "save", "mapping", "naming", "OutputName", "parse_output_name", + "GkylCReader", "GkylReader", "GkylH5Reader", "FlashH5Reader" +] diff --git a/src/postgkyl/io/flash_h5_reader.py b/src/postgkyl/io/flash_h5_reader.py new file mode 100644 index 00000000..ce8a3903 --- /dev/null +++ b/src/postgkyl/io/flash_h5_reader.py @@ -0,0 +1,126 @@ +"""Reader for FLASH code HDF5 output. + +FLASH variable names (for reference, not enforced here): + +- ``dens``: density [g/cc] +- ``tele``/``tion``: electron/ion temperature [K] +- ``velx``/``vely``: fluid velocity [cm/s] +- ``temp``: overall fluid temperature [K] +- ``pres``: pressure [dyn/cm^2] +- ``ye``/``sumy``: used to recover ion/electron density, + ``n_ele = ye * Na * dens``, ``n_ion = sumy * Na * dens`` (``Na`` = Avogadro + number); average ionization ``Z' = ye / sumy``, average atomic mass + ``A' = 1 / sumy``. + +``tables`` (PyTables) is a hard dependency (see ``pyproject.toml``), so this +reader needs no optional-import guard. +""" + +from __future__ import annotations + +import math +from typing import Tuple + +import numpy as np +import tables + +from . import mapping + + +class FlashH5Reader: + """Provides a framework to read FLASH HDF5 output.""" + + def __init__(self, + file_name: str, + ctx: dict | None = None, + var_name: str | None = None, + **kwargs): + """Initialize the instance of the FLASH reader. + + Args: + file_name: path to the ``.h5`` file. + ctx: dict passing context/metadata back to the caller. + var_name: FLASH block variable to read (e.g. ``"dens"``); required by + :meth:`load` but not by :meth:`is_compatible`. + **kwargs: unused; keeps the constructor signature uniform across the + reader registry. + """ + self._file_name = str(file_name) + self.var_name = var_name + + self.ctx = ctx if ctx is not None else {} + + def is_compatible(self) -> bool: + """Checks if the file can be read with the FLASH reader.""" + try: + fh = tables.open_file(self._file_name, "r") + except (tables.exceptions.HDF5ExtError, OSError): + return False + out = "coordinates" in fh.root + fh.close() + return out + + def _read_frame(self) -> tuple: + fh = tables.open_file(self._file_name, "r") + coord = fh.root["coordinates"].read().transpose() + bsize = fh.root["block size"].read().transpose() + ntype = fh.root["node type"].read().transpose() + bdata = fh.root[self.var_name].read().transpose() + fh.close() + + nxb, nyb, _, num_blocks = bdata.shape + res = bsize.min(axis=1) + lower = (coord - bsize / 2).min(axis=1) + upper = (coord + bsize / 2).max(axis=1) + + nxax = math.floor((upper[0] - lower[0]) / (res[0] / nxb)) + nyax = math.floor((upper[1] - lower[1]) / (res[1] / nyb)) + data = np.zeros((nxax, nyax)) + for b in range(num_blocks): + if ntype[b] == 1: + mult = np.ceil(bsize[:, b] / res) + idxx = math.floor( + (coord[0, b] - bsize[0, b] / 2 - lower[0]) / res[0] * nxb) + idxy = math.floor( + (coord[1, b] - bsize[1, b] / 2 - lower[1]) / res[1] * nyb) + for i in range(nxb): + for j in range(nyb): + data[ + idxx + i * int(mult[0]):idxx + (i + 1) * int(mult[0]) + 1, + idxy + j * int(mult[1]):idxy + (j + 1) * int(mult[1]) + 1, + ] = bdata[i, j, 0, b] + return data.shape, lower[:2], upper[:2], data[..., np.newaxis] + + # ---- Exposed functions ----- + def preload(self) -> None: + """Loads metadata. FLASH block reassembly needs the full field, so there + is nothing cheaper to precompute here.""" + + def load(self) -> Tuple[list, np.ndarray]: + """Loads data. + + Returns: + A tuple including a grid list and a data NumPy array. + + Raises: + ValueError: if ``var_name`` was not given. + + Notes: + Needs to be called after ``preload``. + """ + if self.var_name is None: + raise ValueError( + "FlashH5Reader requires 'var_name' (the FLASH block variable to " + "read, e.g. 'dens') to load data.") + + cells, lower, upper, data = self._read_frame() + self.ctx["cells"] = cells + self.ctx["lower"] = lower + self.ctx["upper"] = upper + self.ctx["num_comps"] = data.shape[-1] + self.ctx["grid_type"] = "uniform" + + grid = mapping.uniform_grid(np.asarray(lower, dtype=float), + np.asarray(upper, dtype=float), + np.asarray(cells)) + return grid, data diff --git a/src/postgkyl/io/gkyl_c_reader.py b/src/postgkyl/io/gkyl_c_reader.py new file mode 100644 index 00000000..584c8c19 --- /dev/null +++ b/src/postgkyl/io/gkyl_c_reader.py @@ -0,0 +1,103 @@ +"""``.gkyl`` reading through Gkeyll itself (the primary read path). + +``GkylCReader`` delegates the whole read -- header, grid, allocation, payload, +multi-range stitching -- to ``libg0core.so`` via :mod:`postgkyl.gpython.rio` and +returns the data as a **native** :class:`~postgkyl.gpython.array.GkylArray`, so +modal datasets start life in the modal domain. Python's only jobs are decoding +the msgpack metadata blob into ``ctx`` (same key policy as the pure-Python +reader) and building the NumPy edge grid. + +It declines (``is_compatible() -> False``) when the FFI is unavailable, the +file is not a field file (types 1/3), or a partial load (``axes=``/``comp=``) +was requested -- those fall through to the pure-Python :class:`GkylReader`. +""" + +from __future__ import annotations + +import numpy as np +import msgpack + +from postgkyl import gpython +from . import mapping + + +class GkylCReader: + """Reader protocol implementation backed by ``gkyl_array_rio``.""" + + def __init__(self, + file_name: str, + ctx: dict | None = None, + value_form: str | None = None, + basis_type: str | None = None, + poly_order: int | None = None, + **kwargs): + self.file_name = str(file_name) + self.ctx = ctx if ctx is not None else {} + self._value_form_override = value_form + self._basis_type_override = basis_type + self._poly_order_override = poly_order + # Any partial-load request (axes=, comp=, ...) -> defer to the Python reader. + self._partial = any(v is not None for v in kwargs.get("axes") or ()) or \ + kwargs.get("comp") is not None or \ + bool({k for k in kwargs if k not in ("axes", "comp")}) + + def is_compatible(self) -> bool: + if self._partial or not gpython.available(): + return False + try: + return gpython.rio.file_type( + self.file_name) in gpython.rio.FIELD_FILE_TYPES + except (OSError, RuntimeError): + return False + + def preload(self) -> None: + grid, _, meta, esznc, _ = gpython.rio.read_header(self.file_name) + has_basis = False + if meta: + for key, val in msgpack.unpackb(meta).items(): + if key in ("polyOrder", "poly_order"): + self.ctx["poly_order"] = val + elif key in ("basisType", "basis_type"): + self.ctx["basis_type"] = val + has_basis = True + else: + # Covers "value_form" too, if the writer stamped one directly: + # a file's own metadata is the next-best source of truth once no + # explicit override was given. + self.ctx[key] = val + if self._basis_type_override is not None: + # The writer stamps modal metadata itself; this lets a caller correct + # a missing/mistagged basis_type (e.g. a file with no header metadata + # at all, or one written by a version that mislabeled it) so downstream + # verbs (interpolate, average, integrate, ...) resolve the right basis. + self.ctx["basis_type"] = self._basis_type_override + has_basis = True + if self._poly_order_override is not None: + # Independent of basis_type/value_form: lets a caller correct just the + # polynomial order (e.g. a file with no header metadata at all) without + # asserting anything about modality. + self.ctx["poly_order"] = self._poly_order_override + if has_basis and "value_form" not in self.ctx: + self.ctx["value_form"] = "modal" + if self._value_form_override is not None: + # The writer stamps every field with basis/order metadata, even + # non-DG diagnostic outputs (e.g. a per-cell CFL rate); this lets a + # caller correct a mistagged file to what its values actually are. + # Wins over both the file's own metadata and the "modal" default. + self.ctx["value_form"] = self._value_form_override + self.ctx["cells"] = grid["cells"] + self.ctx["lower"] = grid["lower"] + self.ctx["upper"] = grid["upper"] + self.ctx["num_comps"] = esznc // 8 # payload is float64 + + def load(self): + grid, arr = gpython.rio.read_field(self.file_name) + cells = grid["cells"] + if arr.size != int(np.prod(cells)): + raise IOError( + f"'{self.file_name}': stored cells {arr.size} do not match the " + f"domain {tuple(cells)} (ghost-cell layout?) -- not supported by " + "the Gkeyll read path yet") + edges = mapping.uniform_grid(grid["lower"], grid["upper"], cells) + self.ctx["grid_type"] = "uniform" + return edges, arr diff --git a/src/postgkyl/io/gkyl_h5_reader.py b/src/postgkyl/io/gkyl_h5_reader.py new file mode 100644 index 00000000..aa0cbe6e --- /dev/null +++ b/src/postgkyl/io/gkyl_h5_reader.py @@ -0,0 +1,105 @@ +"""Reader for legacy Gkeyll HDF5 output (predates the native .gkyl binary format). + +``tables`` (PyTables) is a hard dependency (see ``pyproject.toml``), so this +reader needs no optional-import guard. +""" + +from __future__ import annotations + +from typing import Tuple + +import numpy as np +import tables + +from . import mapping + + +class GkylH5Reader: + """Provides a framework to read legacy Gkeyll HDF5 output.""" + + def __init__(self, file_name: str, ctx: dict | None = None, **kwargs): + """Initialize the instance of the legacy Gkeyll HDF5 reader. + + Args: + file_name: path to the ``.h5`` file. + ctx: dict passing context/metadata back to the caller. + **kwargs: unused; keeps the constructor signature uniform across the + reader registry. + """ + self._file_name = str(file_name) + + self.is_frame = False + self.is_diagnostic = False + + self.ctx = ctx if ctx is not None else {} + + def is_compatible(self) -> bool: + """Checks if the file can be read with the legacy Gkeyll HDF5 reader.""" + try: + fh = tables.open_file(self._file_name, "r") + except (tables.exceptions.HDF5ExtError, OSError): + return False + + if "/DataStruct/data" in fh: + self.is_diagnostic = True + if "/StructGridField" in fh: + self.is_frame = True + fh.close() + return self.is_frame or self.is_diagnostic + + def _read_frame(self) -> tuple: + fh = tables.open_file(self._file_name, "r") + + # Postgkyl conventions require the attributes to be arrays even for 1D data. + lower = np.atleast_1d(fh.root.StructGrid._v_attrs.vsLowerBounds) + upper = np.atleast_1d(fh.root.StructGrid._v_attrs.vsUpperBounds) + cells = np.atleast_1d(fh.root.StructGrid._v_attrs.vsNumCells) + if "/timeData" in fh: + self.ctx["time"] = fh.root.timeData._v_attrs.vsTime + + data = fh.root.StructGridField.read() + + fh.close() + return cells, lower, upper, data + + def _read_diagnostic(self) -> tuple: + fh = tables.open_file(self._file_name, "r") + + grid = fh.root.DataStruct.timeMesh.read() + data = fh.root.DataStruct.data.read() + + fh.close() + return [np.squeeze(grid)], [grid[0]], [grid[-1]], data + + # ---- Exposed functions ----- + def preload(self) -> None: + """Loads metadata. Nothing to precompute for this format.""" + + def load(self) -> Tuple[list, np.ndarray]: + """Loads data. + + Returns: + A tuple including a grid list and a data NumPy array. + + Notes: + Needs to be called after ``preload``. + """ + if self.is_frame: + cells, lower, upper, data = self._read_frame() + else: + grid, lower, upper, data = self._read_diagnostic() + cells = grid[0].shape + + self.ctx["cells"] = cells + self.ctx["lower"] = lower + self.ctx["upper"] = upper + self.ctx["num_comps"] = 1 + if len(data.shape) > len(cells): + self.ctx["num_comps"] = data.shape[-1] + + grid = mapping.uniform_grid(np.asarray(lower, dtype=float), + np.asarray(upper, dtype=float), + np.asarray(cells)) + self.ctx["grid_type"] = "uniform" + + return grid, data diff --git a/src/postgkyl/data/gkyl_reader.py b/src/postgkyl/io/gkyl_reader.py similarity index 57% rename from src/postgkyl/data/gkyl_reader.py rename to src/postgkyl/io/gkyl_reader.py index 2e881f99..52951567 100644 --- a/src/postgkyl/data/gkyl_reader.py +++ b/src/postgkyl/io/gkyl_reader.py @@ -1,11 +1,12 @@ """Module including Gkeyll binary reader class.""" -from collections.abc import Iterable from typing import Tuple import msgpack as mp import numpy as np import os.path +from . import mapping + # Format description for raw Gkeyll output file from # gkyl_array_rio_format_desc.h @@ -81,11 +82,15 @@ class GkylReader(object): """Provides a framework to read Gkeyll binary output.""" - def __init__(self, file_name: str, ctx: dict | None = None, - c2p: str = "", c2p_vel: str = "", - axes: tuple | None = (None, None, None, None, None, None), - comp: str | int | None = None, - **kwargs): + def __init__(self, + file_name: str, + ctx: dict | None = None, + axes: tuple | None = (None, None, None, None, None, None), + comp: str | int | None = None, + value_form: str | None = None, + basis_type: str | None = None, + poly_order: int | None = None, + **kwargs): """Initialize the instance of Gkeyll reader. Args: @@ -93,22 +98,29 @@ def __init__(self, file_name: str, ctx: dict | None = None, ctx: dict Passes context variable with metadata. var_name: str = "CartGridField" - c2p: str - Allows to specify a name of the file containing c2p mapping. - c2p_vel: str - Allows to specify a name of the file containing c2p mapping for only the - velocity dimension. axes: tuple Allows to specify the axes to be loaded. comp: int or slice Allows to specify the components to be loaded. + value_form: overrides the ``ctx["value_form"]`` the header + metadata would otherwise imply (e.g. a file tagged with basis + metadata whose stored values are already point values, not modal + coefficients). + basis_type: overrides the ``ctx["basis_type"]`` the header metadata + would otherwise imply -- for files with no basis metadata at all, + or metadata that mislabels the basis actually used. + poly_order: overrides the ``ctx["poly_order"]`` the header metadata + would otherwise imply. Independent of ``basis_type``/``value_form`` -- + it corrects only the polynomial order, asserting nothing about + modality. **kwargs This is not directly used but allowes for unified interface to all the readers we use. """ self.file_name = file_name - self.c2p = c2p - self.c2p_vel = c2p_vel + self._value_form_override = value_form + self._basis_type_override = basis_type + self._poly_order_override = poly_order self.dtf = np.dtype("f8") self.dti = np.dtype("i8") @@ -119,16 +131,15 @@ def __init__(self, file_name: str, ctx: dict | None = None, self.file_type = 1 self.version = 0 - self.lower : np.ndarray - self.upper : np.ndarray - self.num_comps : int - self.cells : np.ndarray + self.lower: np.ndarray + self.upper: np.ndarray + self.num_comps: int + self.cells: np.ndarray if ctx is not None: self.ctx = ctx else: self.ctx = {} - #end if not ("grid_type" in self.ctx.keys()): self.ctx["grid_type"] = "uniform" @@ -141,28 +152,27 @@ def __init__(self, file_name: str, ctx: dict | None = None, if ax is not None: self.partial_load = True self.partial_idxs[i] = str(ax) - #end - #end - #end if comp is not None: self.partial_load = True self.partial_idxs[6] = str(comp) - #end def is_compatible(self) -> bool: """Checks if file can be read with Gkeyll reader.""" try: - magic = np.fromfile(self.file_name, dtype=np.dtype("b"), count=5, offset=0) + magic = np.fromfile(self.file_name, + dtype=np.dtype("b"), + count=5, + offset=0) if np.array_equal(magic, [103, 107, 121, 108, 48]): - self.version = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=5)[0] + self.version = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=5)[0] return True else: return False - #end except: return False - #end - #end # Starting with version 1, .gkyl files contain a header; # Version 0 files only include the real-type info @@ -171,16 +181,26 @@ def _read_header(self) -> None: if self.is_compatible(): self.offset += 5 # Header contatins the gkyl magic sequence - self.version = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + self.version = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 - self.file_type = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + self.file_type = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 - meta_size = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + meta_size = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 # read meta + has_basis = False if meta_size > 0: fh = open(self.file_name, "rb") fh.seek(self.offset) @@ -191,60 +211,91 @@ def _read_header(self) -> None: self.ctx["poly_order"] = unp[key] elif key == "basisType" or key == "basis_type": self.ctx["basis_type"] = unp[key] - self.ctx["is_modal"] = True + has_basis = True else: + # Covers "value_form" too, if the writer stamped one + # directly: a file's own metadata is the next-best source of + # truth once no explicit override was given (see below). self.ctx[key] = unp[key] - #end - #end - #end self.offset += meta_size fh.close() - #end - #end + if self._basis_type_override is not None: + # Wins over the file's own metadata (or its absence): lets a caller + # correct a missing/mislabeled basis_type so downstream verbs + # (interpolate, average, integrate, ...) resolve the right basis. + self.ctx["basis_type"] = self._basis_type_override + has_basis = True + if self._poly_order_override is not None: + # Independent of basis_type/value_form: corrects only the polynomial + # order, asserting nothing about modality. + self.ctx["poly_order"] = self._poly_order_override + if has_basis and "value_form" not in self.ctx: + self.ctx["value_form"] = "modal" + if self._value_form_override is not None: + # Wins over both the file's own metadata and the "modal" default. + self.ctx["value_form"] = self._value_form_override # read real-type - real_type = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + real_type = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] if real_type == 1: self.dtf = np.dtype("f4") self.doffset = 4 - #end self.offset += 8 - #end def _read_t1t3_v1_domain(self) -> None: """Read domain information for file type 1 and 3.""" # read grid dimensions - self.num_dims = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + self.num_dims = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 # read grid shape - self.cells = np.fromfile(self.file_name, dtype=self.dti, count=self.num_dims, offset=self.offset) + self.cells = np.fromfile(self.file_name, + dtype=self.dti, + count=self.num_dims, + offset=self.offset) self.offset += self.num_dims * 8 # read lower/upper - self.lower = np.fromfile(self.file_name, dtype=self.dtf, count=self.num_dims, offset=self.offset) + self.lower = np.fromfile(self.file_name, + dtype=self.dtf, + count=self.num_dims, + offset=self.offset) self.offset += self.num_dims * self.doffset - self.upper = np.fromfile(self.file_name, dtype=self.dtf, count=self.num_dims, offset=self.offset) + self.upper = np.fromfile(self.file_name, + dtype=self.dtf, + count=self.num_dims, + offset=self.offset) self.offset += self.num_dims * self.doffset # read array elem_ez (the div by doffset is as elem_sz includes # sizeof(real_type) = doffset) - elem_sz_raw = int(np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0]) + elem_sz_raw = int( + np.fromfile(self.file_name, dtype=self.dti, count=1, + offset=self.offset)[0]) elem_sz = elem_sz_raw / self.doffset self.num_comps = int(elem_sz) self.offset += 8 # read array size - self.asize = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + self.asize = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 # prep for partial loading - self.orig_size_array = np.zeros(self.num_dims+1, dtype=self.dti) + self.orig_size_array = np.zeros(self.num_dims + 1, dtype=self.dti) self.orig_size_array[:-1] = self.cells.copy() self.orig_size_array[-1] = self.num_comps if self.partial_load: # The offsets are set to zero by default - self.global_offsets = np.zeros((self.num_dims+1, 2), dtype=self.dti) + self.global_offsets = np.zeros((self.num_dims + 1, 2), dtype=self.dti) # The offsets need to be parsed; note that for ":", the Python syntax is used, # i.e., the first index is included, the second is excluded. Negative indices are @@ -262,9 +313,6 @@ def _read_t1t3_v1_domain(self) -> None: self.global_offsets[i, 1] = self.cells[i] - int(stop) elif stop: self.global_offsets[i, 1] = -int(stop) - #end - #end - #end sl = self.partial_idxs[6] if sl.isdigit(): @@ -278,19 +326,17 @@ def _read_t1t3_v1_domain(self) -> None: self.global_offsets[-1, 1] = self.num_comps - int(stop) elif stop: self.global_offsets[-1, 1] = -int(stop) - #end - #end self.cells -= (self.global_offsets[:-1, 1] + self.global_offsets[:-1, 0]) cell_size = (self.upper - self.lower) / self.orig_size_array[:-1] self.lower += self.global_offsets[:-1, 0] * cell_size self.upper -= self.global_offsets[:-1, 1] * cell_size - self.num_comps -= (self.global_offsets[-1, 1] + self.global_offsets[-1, 0]) - #end - #end + self.num_comps -= (self.global_offsets[-1, 1] + + self.global_offsets[-1, 0]) - def _get_block(self, dim : int, out : np.ndarray, idx : int, - dim_offsets : np.ndarray, num_elems : np.ndarray, cells : np.ndarray) -> int: + def _get_block(self, dim: int, out: np.ndarray, idx: int, + dim_offsets: np.ndarray, num_elems: np.ndarray, + cells: np.ndarray) -> int: """Reads a block of data. A recursion is used to read the data from the fastest going index (the last one; @@ -298,122 +344,143 @@ def _get_block(self, dim : int, out : np.ndarray, idx : int, """ if dim == self.num_dims: self.offset += dim_offsets[-1, 0] * self.doffset - out[idx : idx+self.num_comps] = np.fromfile(file=self.file_name, - dtype=self.dtf, count=self.num_comps, offset=self.offset) + out[idx:idx + self.num_comps] = np.fromfile(file=self.file_name, + dtype=self.dtf, + count=self.num_comps, + offset=self.offset) self.offset += (self.num_comps + dim_offsets[-1, 1]) * self.doffset idx += self.num_comps else: - self.offset += dim_offsets[dim, 0] * np.prod(num_elems[dim+1:]) * self.doffset + self.offset += dim_offsets[dim, 0] * np.prod( + num_elems[dim + 1:]) * self.doffset for _ in range(cells[dim]): - idx = self._get_block(dim=dim+1, out=out, idx=idx, dim_offsets=dim_offsets, - num_elems=num_elems, cells=cells) - #end - self.offset += dim_offsets[dim, 1] * np.prod(num_elems[dim+1:]) * self.doffset - #end + idx = self._get_block(dim=dim + 1, + out=out, + idx=idx, + dim_offsets=dim_offsets, + num_elems=num_elems, + cells=cells) + self.offset += dim_offsets[dim, 1] * np.prod( + num_elems[dim + 1:]) * self.doffset return idx - #end - def _get_data(self, count : int, - lo_idx : np.ndarray | None = None, up_idx : np.ndarray | None = None) -> Tuple[np.ndarray, Tuple]: + def _get_data(self, + count: int, + lo_idx: np.ndarray | None = None, + up_idx: np.ndarray | None = None) -> Tuple[np.ndarray, Tuple]: """Read raw data and account for partial load.""" slices = [] gshape = np.ones(self.num_dims + 1, dtype=self.dti) gshape[-1] = self.num_comps if not self.partial_load: - out = np.fromfile(self.file_name, dtype=self.dtf, count=count, offset=self.offset) + out = np.fromfile(self.file_name, + dtype=self.dtf, + count=count, + offset=self.offset) self.offset += count * self.doffset if lo_idx is not None: for d in range(self.num_dims): gshape[d] = up_idx[d] - lo_idx[d] + 1 - #end - slices = [slice(lo_idx[d] - 1, up_idx[d]) for d in range(self.num_dims)] # Gkeyll is 1-indexed + slices = [ + slice(lo_idx[d] - 1, up_idx[d]) for d in range(self.num_dims) + ] # Gkeyll is 1-indexed else: for d in range(self.num_dims): gshape[d] = self.cells[d] - #end - #end else: if lo_idx is None: - lo_idx = np.ones(self.num_dims, dtype=self.dti) # Gkeyll index is 1-indexed - #end + lo_idx = np.ones(self.num_dims, + dtype=self.dti) # Gkeyll index is 1-indexed if up_idx is None: up_idx = self.orig_size_array[:-1] - #end num_elems = self.orig_size_array.copy() # Adjust the offsets for the partial load for distributed memory data dim_offsets = np.zeros_like(self.global_offsets, dtype=self.dti) dim_offsets[:-1, 0] = self.global_offsets[:-1, 0] - (lo_idx - 1) - dim_offsets[:-1, 1] = self.global_offsets[:-1, 1] - (num_elems[:-1] - up_idx) + dim_offsets[:-1, + 1] = self.global_offsets[:-1, 1] - (num_elems[:-1] - up_idx) dim_offsets[-1, :] = self.global_offsets[-1, :] dim_offsets = dim_offsets.clip(min=0) # Calculate the size to allocate the memory - num_elems[:-1] = up_idx - lo_idx + 1 # Gkeyll index is 1-indexed + num_elems[:-1] = up_idx - lo_idx + 1 # Gkeyll index is 1-indexed cells = num_elems[:-1] - dim_offsets[:-1, 1] - dim_offsets[:-1, 0] if np.any(cells < 1): self.offset += count * self.doffset return np.array([]), tuple(slices) - #end size = np.prod(cells) * self.num_comps - out = np.zeros(size, dtype=self.dtf) # Allocate space for the data - self._get_block(dim=0, out=out, idx=0, dim_offsets=dim_offsets, - num_elems=num_elems, cells=cells) + out = np.zeros(size, dtype=self.dtf) # Allocate space for the data + self._get_block(dim=0, + out=out, + idx=0, + dim_offsets=dim_offsets, + num_elems=num_elems, + cells=cells) lo_idx = (lo_idx - self.global_offsets[:-1, 0]).clip(min=1) - up_idx = (up_idx - self.global_offsets[:-1, 0] - dim_offsets[:-1, 1]).clip(min=1) + up_idx = (up_idx - self.global_offsets[:-1, 0] - + dim_offsets[:-1, 1]).clip(min=1) for d in range(self.num_dims): gshape[d] = up_idx[d] - lo_idx[d] + 1 - #end - slices = [slice(lo_idx[d] - 1, up_idx[d]) for d in range(self.num_dims)] # Gkeyll is 1-indexed - #end + slices = [slice(lo_idx[d] - 1, up_idx[d]) + for d in range(self.num_dims)] # Gkeyll is 1-indexed return out.reshape(gshape, order="C"), tuple(slices) - #end def _read_t1_v1_data(self) -> np.ndarray: - """Reat field data for file type 1.""" - data, _ = self._get_data(self.asize*self.num_comps) + """Read field data for file type 1.""" + data, _ = self._get_data(self.asize * self.num_comps) return data def _read_t3_v1_data(self) -> np.ndarray: """Read field data for file type 3.""" # get the number of stored ranges - num_range = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + num_range = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 gshape = np.ones(self.num_dims + 1, dtype=self.dti) for d in range(self.num_dims): gshape[d] = self.cells[d] - #end gshape[-1] = self.num_comps - data = np.zeros(gshape, dtype=self.dtf) # Allocate space for the data + data = np.zeros(gshape, dtype=self.dtf) # Allocate space for the data for _ in range(num_range): - lo_idx = np.fromfile(self.file_name, dtype=self.dti, count=self.num_dims, offset=self.offset) + lo_idx = np.fromfile(self.file_name, + dtype=self.dti, + count=self.num_dims, + offset=self.offset) self.offset += self.num_dims * 8 - up_idx = np.fromfile(self.file_name, dtype=self.dti, count=self.num_dims, offset=self.offset) + up_idx = np.fromfile(self.file_name, + dtype=self.dti, + count=self.num_dims, + offset=self.offset) self.offset += self.num_dims * 8 - asize = np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0] + asize = np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0] self.offset += 8 #data_raw = np.fromfile(self.file_name, dtype=self.dtf, count=asize*self.num_comps, # offset=self.offset) #self.offset += asize * self.num_comps * self.doffset - data_block, slices = self._get_data(count=asize*self.orig_size_array[-1], - lo_idx=lo_idx, up_idx=up_idx) + data_block, slices = self._get_data(count=asize * + self.orig_size_array[-1], + lo_idx=lo_idx, + up_idx=up_idx) if len(data_block) == 0: continue - #end data[slices] = data_block - #end return data - #end def _read_t2_v1(self) -> Tuple[list, np.ndarray]: """Read dynvector data for file type 2.""" @@ -421,18 +488,31 @@ def _read_t2_v1(self) -> Tuple[list, np.ndarray]: time = np.array([]) data = np.array([[]]) while True: # Python does not have DO .. WHILE loop - elem_sz_raw = int(np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0]) + elem_sz_raw = int( + np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0]) num_comps = int(elem_sz_raw / self.doffset) self.offset += 8 - loop_cells = int(np.fromfile(self.file_name, dtype=self.dti, count=1, offset=self.offset)[0]) + loop_cells = int( + np.fromfile(self.file_name, + dtype=self.dti, + count=1, + offset=self.offset)[0]) self.offset += 8 - loop_time = np.fromfile(self.file_name, dtype=self.dtf, count=loop_cells, offset=self.offset) + loop_time = np.fromfile(self.file_name, + dtype=self.dtf, + count=loop_cells, + offset=self.offset) self.offset += loop_cells * 8 - data_raw = np.fromfile(self.file_name, dtype=self.dtf, count=num_comps * loop_cells, - offset=self.offset) + data_raw = np.fromfile(self.file_name, + dtype=self.dtf, + count=num_comps * loop_cells, + offset=self.offset) self.offset += loop_cells * elem_sz_raw gshape = np.array((loop_cells, num_comps), dtype=self.dti) @@ -441,21 +521,16 @@ def _read_t2_v1(self) -> Tuple[list, np.ndarray]: data = data_raw.reshape(gshape, order="C") else: data = np.append(data, data_raw.reshape(gshape, order="C"), axis=0) - #end cells += loop_cells if self.offset >= os.path.getsize(self.file_name): break - #end self._read_header() if self.file_type != 2: raise TypeError("Inconsitent data in g0 dynVector file.") - #end - #end self.cells = [cells] self.lower = np.atleast_1d(time.min()) self.upper = np.atleast_1d(time.max()) return time, data - #end # ---- Exposed functions ----- def preload(self) -> None: @@ -468,9 +543,6 @@ def preload(self) -> None: self.ctx["lower"] = self.lower self.ctx["upper"] = self.upper self.ctx["num_comps"] = self.num_comps - #end - #end - #end def load(self) -> Tuple[list, np.ndarray]: """Loads data. @@ -490,71 +562,17 @@ def load(self) -> Tuple[list, np.ndarray]: data = self._read_t3_v1_data() else: raise TypeError("This g0 format is not presently supported") - #end # Load or construct grid - num_dims = len(self.cells) if time is not None: grid = [time] if self.ctx: self.ctx["grid_type"] = "nodal" - #end - elif self.c2p: - grid_reader = GkylReader(self.c2p) - grid_reader.preload() - _, tmp = grid_reader.load() - num_comps = tmp.shape[-1] - num_coeff = num_comps / num_dims - grid = [tmp[..., int(d * num_coeff) : int((d + 1)*num_coeff)] for d in range(num_dims)] - if self.ctx: - self.ctx["grid_type"] = "c2p" - #end - elif self.c2p_vel: - grid_reader = GkylReader(self.c2p_vel) - grid_reader.preload() - _, tmp = grid_reader.load() - - num_vdim = len(tmp.shape) - 1 - num_cdim = num_dims - num_vdim - if self.ctx: - self.ctx["num_vdim"] = num_vdim - self.ctx["num_cdim"] = num_cdim - #end - - # Create uniform configuration space grid - grid = [np.linspace(self.lower[d], self.upper[d], self.cells[d] + 1) for d in range(num_cdim)] - - # Create non-uniform velocity grid - num_comps = tmp.shape[-1] - num_coeff = num_comps / num_vdim - for d in range(num_vdim): - idx = [0] * (num_vdim + 1) - idx[d] = slice(None) - idx[-1] = slice(int(d * num_coeff), int((d + 1) * num_coeff)) - grid.append(tmp[tuple(idx)]) - #end - - if self.ctx: - self.ctx["grid_type"] = "c2p_vel" - #end else: # Create sparse unifrom grid - # Adjust for ghost cells - dz = (self.upper - self.lower) / self.cells - for d in range(num_dims): - if self.cells[d] != data.shape[d]: - ngl = int(np.floor((self.cells[d] - data.shape[d]) * 0.5)) - ngu = int(np.ceil((self.cells[d] - data.shape[d]) * 0.5)) - self.cells[d] = data.shape[d] - self.lower[d] = self.lower[d] - ngl * dz[d] - self.upper[d] = self.upper[d] + ngu * dz[d] - #end - #end - grid = [np.linspace(self.lower[d], self.upper[d], self.cells[d] + 1) for d in range(num_dims)] + mapping.adjust_for_ghost_cells(self.lower, self.upper, self.cells, + data.shape) + grid = mapping.uniform_grid(self.lower, self.upper, self.cells) if self.ctx: self.ctx["grid_type"] = "uniform" - #end - #end return grid, data - #end -#end diff --git a/src/postgkyl/io/mapping.py b/src/postgkyl/io/mapping.py new file mode 100644 index 00000000..bb755085 --- /dev/null +++ b/src/postgkyl/io/mapping.py @@ -0,0 +1,61 @@ +"""Read-time grid construction for Gkeyll output. + +A Gkeyll field stores only its *values*; at read time the grid is built +uniformly from the stored bounds (corrected for ghost cells). Coordinate +(computational-to-physical) mappings are *not* applied while reading -- they +are applied afterwards, on already-loaded data, by the ``map`` verb +(``operations/map.py``, backed by ``dg/map.py``; see ``MAPPING.md``). + +``uniform_grid``/``adjust_for_ghost_cells`` build the read-time uniform grid. +``c2p_grid`` is unused by ``operations/map.py`` (which evaluates a mapping's DG +coefficients directly via ``dg.map_grid``/``gpython.basis.eval_matrix`` +rather than splitting packed node coordinates) but is kept for any reader +that still needs to split a mapping field's packed per-dimension node block. +""" + +from __future__ import annotations + +import numpy as np + + +def adjust_for_ghost_cells(lower: np.ndarray, upper: np.ndarray, + cells: np.ndarray, data_shape: tuple) -> tuple: + """Shrink the cell count / extend the bounds to account for ghost cells. + + When the stored data has fewer cells along a dimension than ``cells`` + advertises, the difference is ghost cells; the bounds are pushed out by the + ghost-cell width so the resulting grid still maps onto the data. ``lower``, + ``upper`` and ``cells`` are mutated in place and also returned. + """ + num_dims = len(cells) + dz = (upper - lower) / cells + for d in range(num_dims): + if cells[d] != data_shape[d]: + ngl = int(np.floor((cells[d] - data_shape[d]) * 0.5)) + ngu = int(np.ceil((cells[d] - data_shape[d]) * 0.5)) + cells[d] = data_shape[d] + lower[d] = lower[d] - ngl * dz[d] + upper[d] = upper[d] + ngu * dz[d] + return lower, upper, cells + + +def uniform_grid(lower: np.ndarray, upper: np.ndarray, + cells: np.ndarray) -> list: + """A uniform nodal grid: ``cells[d] + 1`` edges per dimension.""" + return [ + np.linspace(lower[d], upper[d], cells[d] + 1) for d in range(len(cells)) + ] + + +def c2p_grid(nodes: np.ndarray, num_dims: int) -> list: + """Split a ``mapc2p`` node array into a per-dimension block of coefficients. + + The mapping file packs every dimension's node coordinates on the last axis; + this slices that axis into ``num_dims`` equal blocks. + """ + num_comps = nodes.shape[-1] + num_coeff = num_comps / num_dims + return [ + nodes[..., int(d * num_coeff):int((d + 1) * num_coeff)] + for d in range(num_dims) + ] diff --git a/src/postgkyl/io/naming.py b/src/postgkyl/io/naming.py new file mode 100644 index 00000000..de57314f --- /dev/null +++ b/src/postgkyl/io/naming.py @@ -0,0 +1,157 @@ +"""Gkeyll's output-file naming convention -- the ONE home for reading a +dataset's *identity* out of its path. + +A Gkeyll output file name encodes four facts:: + + rt_gk_multib_sheath_1x2v_p1_b2-geo_int_B3.gkyl + `------------ sim ---------' `bl' `-quantity-' + + gk_lorentzian_mirror-elc_M0_1.gkyl + `------- sim -------' `-quan-' frame + +- **sim** the simulation name (everything before the last ``'-'``), +- **block** the multiblock block index, the ``_b`` suffix of the sim + part; ``None`` for a single-block run, +- **quantity** the output name (species/moment/geometry field), +- **frame** the trailing ``_`` of the quantity, when present. + +Before this module the convention was re-derived in three places with three +slightly different rules (``diagnostics.gk.rz._file_prefix``'s +``rsplit('-', 1)``, ``diagnostics.discovery.find_output_stems``'s +``_\\d+$`` strip, and the animation operation's port of main's +``utils.set_frame``, which recovered a frame index by diffing the loaded file +names character by character). It lives in ``io`` because it is knowledge +about Gkeyll's *files*, which is what this layer owns, and because ``io`` is +below every consumer (``gdatastate`` stamps it into ``ctx`` at load time; +``diagnostics`` builds directory discovery on top of it). + +The parser is *pure*: it never touches the filesystem. ``os.path.exists`` is +the caller's business. +""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass + +# The multiblock block index: Gkeyll writes "_b-.gkyl" +# (see the ``name + '_b*-'`` prefix built by ``diagnostics.gk.nodes`` +# and main's ``nodes``). Digits are required, so a simulation legitimately +# named e.g. "gk_beta_scan" is never mistaken for block "eta_scan". +_BLOCK_RE = re.compile(r"^(?P.*)_b(?P\d+)$") + +# The frame index: a trailing "_" run. Requiring *only* digits is what +# keeps a geometry field like "geo_int_B3" (frame-less) from being read as +# quantity "geo_int_B" at frame 3. +_FRAME_RE = re.compile(r"^(?P.*)_(?P\d+)$") + +_RESTART_SUFFIX = "_restart" + + +@dataclass(frozen=True) +class OutputName: + """The identity of one Gkeyll output file, parsed from its path. + + Attributes: + directory: The path's directory part ("" for a bare file name). + sim: The simulation name, with any ``_b`` block suffix removed. + block: The multiblock block index, or ``None`` for single-block output. + quantity: The output name, with any frame index and ``_restart`` removed. + frame: The frame index, or ``None`` when the name carries none. + restart: True when the name carried a ``_restart`` suffix. + """ + + directory: str + sim: str + block: int | None + quantity: str + frame: int | None + restart: bool = False + + @property + def prefix(self) -> str: + """The ``'/[_b]'`` path every sibling file of this *block* + shares -- what a geometry lookup appends ``'-geo_int_nodes.gkyl'`` to. + + For single-block output this is exactly the old + ``rz._file_prefix`` (the part of the path before the last ``'-'``). + """ + base = self.sim if self.block is None else f"{self.sim}_b{self.block:d}" + return os.path.join(self.directory, base) if self.directory else base + + @property + def stem(self) -> str: + """``'[_b]-'`` -- the file name with directory, frame + index, ``_restart`` and extension stripped (``discovery``'s notion of a + stem).""" + tail = f"-{self.quantity}" if self.quantity else "" + return f"{os.path.basename(self.prefix)}{tail}" + + @property + def field_key(self) -> tuple: + """What two files of the **same field on different blocks** share. + + Deliberately excludes ``block`` (and the directory): it is the key + :func:`postgkyl.gdatastate.collection.group_blocks` partitions a working + set on. + """ + return (self.sim, self.quantity, self.frame) + + +def parse_output_name(path: str | None) -> OutputName | None: + """Parse a Gkeyll output path into its :class:`OutputName` identity. + + Args: + path: A file path, e.g. ``"data/sim_b2-elc_M0_7.gkyl"``. May be any + extension, or none. + + Returns: + The parsed identity, or ``None`` for an empty/absent path (a dataset a + verb computed rather than read from disk). + + Notes: + The split is deliberately total -- every non-empty path parses. A name + with no ``'-'`` at all (out of convention) yields the whole stem as + ``sim`` and an empty ``quantity``; the frame index is then taken off the + ``sim``, since that is the only component there is. + """ + if not path: + return None + directory, base = os.path.split(str(path)) + stem = os.path.splitext(base)[0] + + restart = stem.endswith(_RESTART_SUFFIX) + if restart: + stem = stem[:-len(_RESTART_SUFFIX)] + + if "-" in stem: + sim_part, quantity = stem.rsplit("-", 1) + else: + sim_part, quantity = stem, "" + + # The frame index trails the *last* component of the name -- the quantity + # normally, the sim itself for a dash-less name. + tail = quantity or sim_part + frame = None + match = _FRAME_RE.match(tail) + if match: + tail = match.group("quantity") + frame = int(match.group("frame")) + if quantity: + quantity = tail + else: + sim_part = tail + + block = None + match = _BLOCK_RE.match(sim_part) + if match: + sim_part = match.group("sim") + block = int(match.group("block")) + + return OutputName(directory=directory, + sim=sim_part, + block=block, + quantity=quantity, + frame=frame, + restart=restart) diff --git a/src/postgkyl/io/writer.py b/src/postgkyl/io/writer.py new file mode 100644 index 00000000..9701b122 --- /dev/null +++ b/src/postgkyl/io/writer.py @@ -0,0 +1,252 @@ +"""Write a dataset back to disk. + +A leaf module: it consumes the read-only *surface* of a dataset (the same +properties the readers fill) and never imports ``gdatastate``/``operations``. Supports the +Gkeyll binary ``.gkyl`` format (round-trips with :class:`GkylReader`), plain +ASCII ``.txt``, NumPy ``.npy``, and legacy VTK structured-grid ``.vtk`` +(for external 3-D/VR viewers such as ParaView). +""" + +from __future__ import annotations + +import json +import os +import re +from typing import Literal, Protocol + +import msgpack +import numpy as np +import pyvista as pv + +from postgkyl.cli_spec import ( + CommandSpec, + Execution, + ResultPolicy, + Section, + command, +) +from postgkyl.numerics import nodal_to_cell_centered_grid + +# ctx keys that are either structural (already carried by the binary header +# itself, e.g. cells/lower/upper) or postgkyl's own session-only bookkeeping +# (recomputed by the reader from the meta below) -- never part of the +# msgpack meta blob Gkeyll writes. +_INTERNAL_CTX_KEYS = frozenset({ + "cells", + "lower", + "upper", + "num_comps", + "num_dims", + "grid_type", + "value_form", + "num_quad", + "interpolated", + "var_names", + # The source file's parsed identity (see io.naming): derived from the + # *path*, never stored in the file. Writing it would be actively wrong -- + # reloading the output under a different name would then find a stale + # sim/block in the header, and GDataState's setdefault lets header + # metadata win over the parsed name. "frame" is deliberately NOT here: + # Gkeyll itself writes that one. + "sim", + "block", + "quantity", +}) + +# ctx uses postgkyl's snake_case names; Gkeyll's own meta blob (and anything +# else that reads the file) expects the original camelCase keys. +_CTX_TO_META_KEY = {"poly_order": "polyOrder", "basis_type": "basisType"} + + +class _WritableDataset(Protocol): + """Read-only dataset surface consumed by the format writers.""" + + num_dims: int + num_comps: int + num_cells: np.ndarray + bounds: tuple[np.ndarray, np.ndarray] + values: object + grid: list + ctx: dict + + +@command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_EACH, + result=ResultPolicy.VALUE)) +def save(data: _WritableDataset, + out_name: str = "", + extension: Literal["gkyl", "txt", "npy", "vtk"] = "gkyl", + var_name: str = "CartGridField") -> str: + """Write ``data`` to ``out_name`` in the requested ``extension``. + + Args: + data: a dataset exposing ``num_dims``/``num_comps``/``num_cells``/ + ``bounds``/``values``/``grid``/``ctx`` (a ``GDataState`` or subclass). + out_name: output path; when empty a name is derived from the source file. + extension: one of ``"gkyl"`` (default), ``"txt"``, ``"npy"``, ``"vtk"``. + var_name: unused placeholder kept for interface symmetry. + + Returns: + The path actually written. + """ + if not out_name: + src = getattr(data, "_file_name", "") or "" + stem = src.split(".", maxsplit=1)[0].strip("_") if src else "gdata" + out_name = f"{stem}_mod.{extension}" + elif out_name.split(".")[-1] != extension: + out_name += "." + extension + + num_dims = data.num_dims + num_comps = data.num_comps + num_cells = data.num_cells + lo, up = data.bounds + values = data.values + + if extension == "gkyl": + ctx = getattr(data, "ctx", {}) or {} + _write_gkyl(out_name, num_dims, num_comps, num_cells, lo, up, values, ctx) + elif extension == "npy": + np.save(out_name, np.asarray(values).squeeze()) + elif extension == "txt": + _write_txt(out_name, data, num_dims, num_comps, num_cells, values) + elif extension == "vtk": + _write_vtk(out_name, data, num_dims, num_cells, values) + else: + raise ValueError(f"Unsupported write extension '{extension}'") + return out_name + + +def _build_meta(ctx: dict) -> dict: + """Translate ``ctx`` back into the msgpack meta blob Gkeyll itself writes + (poly order, basis type, physical params, time/frame stamps, ...) -- + everything except the structural/session-only keys in + ``_INTERNAL_CTX_KEYS``.""" + meta = {} + for key, val in ctx.items(): + if key in _INTERNAL_CTX_KEYS: + continue + meta[_CTX_TO_META_KEY.get(key, key)] = _to_msgpack_safe(val) + return meta + + +def _to_msgpack_safe(val): + if isinstance(val, np.generic): + return val.item() + if isinstance(val, np.ndarray): + return val.tolist() + return val + + +def _write_gkyl(out_name, num_dims, num_comps, num_cells, lo, up, values, + ctx) -> None: + dti = np.dtype("i8") + dtf = np.dtype("f8") + meta = _build_meta(ctx) + packed = msgpack.packb(meta, use_bin_type=True) if meta else b"" + with open(out_name, "wb") as fh: + np.array([103, 107, 121, 108, 48], + dtype=np.dtype("b")).tofile(fh, sep="") # 'gkyl0' + np.array([1], dtype=dti).tofile(fh, sep="") # version 1 + np.array([1], dtype=dti).tofile(fh, sep="") # file type 1 (field) + np.array([len(packed)], dtype=dti).tofile(fh, sep="") # meta size + if packed: + fh.write(packed) + np.array([2], dtype=dti).tofile(fh, sep="") # real type (f8) + np.array([num_dims], dtype=dti).tofile(fh, sep="") + np.array(num_cells, dtype=dti).tofile(fh, sep="") + np.array(lo, dtype=dtf).tofile(fh, sep="") + np.array(up, dtype=dtf).tofile(fh, sep="") + np.array([num_comps * 8], dtype=dti).tofile(fh, sep="") # elem_sz + np.array([int(np.prod(num_cells))], dtype=dti).tofile(fh, sep="") # asize + np.array(values, dtype=dtf).tofile(fh, sep="") + + +def _write_txt(out_name, data, num_dims, num_comps, num_cells, values) -> None: + grid = [0.5 * (g[1:] + g[:-1]) for g in data.grid] # cell centers + num_rows = int(np.prod(num_cells)) + basis = np.full(num_dims, 1.0) + for d in range(num_dims - 1): + basis[d] = np.prod(num_cells[(d + 1):]) + with open(out_name, "w", encoding="utf-8") as fh: + for i in range(num_rows): + idx = i + idxs = np.zeros(num_dims, np.int32) + for d in range(num_dims): + idxs[d] = int(idx // basis[d]) + idx = idx % basis[d] + cells = [f"{grid[d][idxs[d]]:.15e}" for d in range(num_dims)] + comps = [f"{values[tuple(idxs)][c]:.15e}" for c in range(num_comps)] + fh.write(", ".join(cells + comps) + "\n") + + +def _write_vtk(out_name, data, num_dims, num_cells, values) -> None: + """Write a legacy VTK structured-grid file via PyVista. + + 1-D/2-D fields are written as a height-mapped surface (the field value + becomes the missing coordinate, e.g. z for a 1-D line); 3-D fields are + written as a volume with the field stored as point data ``"f_raw"``. + """ + if num_dims not in (1, 2, 3): + raise ValueError(f"VTK output supports 1-3 dimensions, got {num_dims}") + + n_grid = nodal_to_cell_centered_grid(data.grid, num_cells, meshgrid=True) + fval = np.asarray(values).squeeze() + if num_dims == 1: + x = n_grid[0] + y = np.zeros_like(x) + z = fval + elif num_dims == 2: + x, y = n_grid + z = fval + else: + x, y, z = n_grid + + grid3d = pv.StructuredGrid(x, y, z) + grid3d["f_raw"] = fval.ravel(order="F") + grid3d.save(out_name) + _update_vtk_series_file(data, out_name) + + +def _update_vtk_series_file(data, out_name: str) -> None: + """Create or update ParaView ``.series`` metadata for VTK file-series + time playback: each write of a frame-numbered file appends (or refreshes) + its entry, keyed by the series' shared stem.""" + out_dir = os.path.dirname(out_name) + out_file = os.path.basename(out_name) + stem, ext = os.path.splitext(out_file) + match = re.match(r"^(.*?)(?:[_-]?(\d+))$", stem) + if match and match.group(1): + series_stem = match.group(1).rstrip("_-") or stem + else: + series_stem = stem + + series_path = os.path.join(out_dir, f"{series_stem}{ext}.series") + time_value = float(data.ctx.get("time", data.ctx.get("frame", 0.0))) + rel_file = os.path.relpath(out_name, out_dir if out_dir else ".") + + series_data = {"file-series-version": "1.0", "files": []} + if os.path.exists(series_path): + try: + with open(series_path, "r", encoding="utf-8") as fh: + loaded = json.load(fh) + if isinstance(loaded, dict) and isinstance(loaded.get("files"), list): + series_data = loaded + series_data.setdefault("file-series-version", "1.0") + except (OSError, json.JSONDecodeError): + pass + + replaced = False + for entry in series_data["files"]: + if entry.get("name") == rel_file: + entry["time"] = time_value + replaced = True + break + if not replaced: + series_data["files"].append({"name": rel_file, "time": time_value}) + + series_data["files"].sort( + key=lambda x: (float(x.get("time", 0.0)), x.get("name", ""))) + with open(series_path, "w", encoding="utf-8") as fh: + json.dump(series_data, fh, indent=2) + fh.write("\n") diff --git a/src/postgkyl/modalDG/__init__.py b/src/postgkyl/modalDG/__init__.py deleted file mode 100644 index 2acf2e27..00000000 --- a/src/postgkyl/modalDG/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .interpolate import interpolate - -from . import kernels diff --git a/src/postgkyl/modalDG/interpolate.py b/src/postgkyl/modalDG/interpolate.py deleted file mode 100644 index bb57d8ea..00000000 --- a/src/postgkyl/modalDG/interpolate.py +++ /dev/null @@ -1,134 +0,0 @@ -import numpy as np -# from postgkyl.data.data import Data - -from postgkyl.modalDG.kernels import expand_1d, expand_2d, expand_3d, expand_4d, expand_5d, expand_6d - - -def interpolate(data, poly_order=None, nodes=None, externalGrid=None): - if poly_order is None and data.poly_order is not None: - poly_order = data.poly_order - else: - # Something bad happened :D - pass - # end - - # Read grid information from input file. - num_dims = data.get_num_dims() - lower, upper = data.get_bounds() - numCells = data.get_num_cells() - - # If user specifies an interpolation grid, use it. Otherwise calculate. - if externalGrid: - intGrid = externalGrid - else: - intGrid = [ - np.linspace(lower[d], upper[d], numCells[d] * (poly_order + 1) + 1) - for d in range(num_dims) - ] - # end - - # Calculate interpolation nodes for each element. - if not nodes: - dx = 2 / (poly_order + 1) - nodes = np.linspace(-1 + dx / 2, 1 - dx / 2, poly_order + 1) - # end - - # Set up array for interp node values - values = data.get_values() - intValues = np.zeros(np.int32(numCells * len(nodes))) - intValues = intValues[..., np.newaxis] - - # Iterating through the node list, calculate value at each node for each element - # simultaneously, one dimension at a time. - # TODO: Rework for num_dims > 3, currently very slow. - if num_dims == 1: - for i, x in enumerate(nodes): - intValues[i :: len(nodes), 0] = expand_1d[int(poly_order - 1)](values, x) - # end - - elif num_dims == 2: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - intValues[i :: len(nodes), j :: len(nodes), 0] = expand_2d[int(poly_order - 1)]( - values, x, y - ) - # end - # end - # end - - elif num_dims == 3: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - intValues[i :: len(nodes), j :: len(nodes), k :: len(nodes), 0] = expand_3d[ - int(poly_order - 1) - ](values, x, y, z) - # end - # end - # end - # end - - elif num_dims == 4: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - for l, r in enumerate(nodes): - intValues[ - i :: len(nodes), j :: len(nodes), k :: len(nodes), l :: len(nodes), 0 - ] = expand_4d[int(poly_order - 1)](values, x, y, z, r) - # end - # end - # end - # end - # end - - elif num_dims == 5: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - for l, r in enumerate(nodes): - for m, s in enumerate(nodes): - intValues[ - i :: len(nodes), - j :: len(nodes), - k :: len(nodes), - l :: len(nodes), - m :: len(nodes), - 0, - ] = expand_5d[int(poly_order - 1)](values, x, y, z, r, s) - # end - # end - # end - # end - # end - # end - - elif num_dims == 6: - for i, x in enumerate(nodes): - for j, y in enumerate(nodes): - for k, z in enumerate(nodes): - for l, r in enumerate(nodes): - for m, s in enumerate(nodes): - for n, t in enumerate(nodes): - intValues[ - i :: len(nodes), - j :: len(nodes), - k :: len(nodes), - l :: len(nodes), - m :: len(nodes), - n :: len(nodes), - 0, - ] = expand_6d[int(poly_order - 1)](values, x, y, z, r, s, t) - # end - # end - # end - # end - # end - # end - - # Hardcoded stack - data.pushGrid(intGrid) - data.pushValues(intValues) - - -# end diff --git a/src/postgkyl/modalDG/kernels/__init__.py b/src/postgkyl/modalDG/kernels/__init__.py deleted file mode 100644 index dddacfc2..00000000 --- a/src/postgkyl/modalDG/kernels/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .expand1d import expand_1d -from .expand2d import expand_2d -from .expand3d import expand_3d -from .expand4d import expand_4d -from .expand5d import expand_5d -from .expand6d import expand_6d diff --git a/src/postgkyl/modalDG/kernels/expand1d.py b/src/postgkyl/modalDG/kernels/expand1d.py deleted file mode 100644 index 7ed0f1fa..00000000 --- a/src/postgkyl/modalDG/kernels/expand1d.py +++ /dev/null @@ -1,45 +0,0 @@ -def _expand_1d1p(f, x): - return 1.224744871391589 * f[..., 1] * x + 0.7071067811865475 * f[..., 0] - - -# end - - -def _expand_1d2p(f, x): - return ( - 2.371708245126284 * f[..., 2] * (x ** 2 - 0.3333333333333333) - + 1.224744871391589 * f[..., 1] * x - + 0.7071067811865475 * f[..., 0] - ) - - -# end - - -def _expand_1d3p(f, x): - return ( - 4.677071733467426 * f[..., 3] * (x ** 3 - 0.6 * x) - + 2.371708245126284 * f[..., 2] * (x ** 2 - 0.3333333333333333) - + 1.224744871391589 * f[..., 1] * x - + 0.7071067811865475 * f[..., 0] - ) - - -# end - - -def _expand_1d4p(f, x): - return ( - 9.280776503073433 - * f[..., 4] - * (x ** 4 - 0.8571428571428571 * (x ** 2 - 0.3333333333333333) - 0.2) - + 4.677071733467426 * f[..., 3] * (x ** 3 - 0.6 * x) - + 2.371708245126284 * f[..., 2] * (x ** 2 - 0.3333333333333333) - + 1.224744871391589 * f[..., 1] * x - + 0.7071067811865475 * f[..., 0] - ) - - -# end - -expand_1d = [_expand_1d1p, _expand_1d2p, _expand_1d3p, _expand_1d4p] diff --git a/src/postgkyl/modalDG/kernels/expand2d.py b/src/postgkyl/modalDG/kernels/expand2d.py deleted file mode 100755 index 21e4b386..00000000 --- a/src/postgkyl/modalDG/kernels/expand2d.py +++ /dev/null @@ -1,86 +0,0 @@ -def _expand_2d1p(f, x, y): - return ( - 1.5 * f[..., 3] * x * y - + 0.8660254037844386 * f[..., 2] * y - + 0.8660254037844386 * f[..., 1] * x - + 0.5 * f[..., 0] - ) - - -# end - - -def _expand_2d2p(f, x, y): - return ( - 2.904737509655563 * f[..., 7] * (x * y**2 - 0.3333333333333333 * x) - + 1.677050983124842 * f[..., 5] * (y**2 - 0.3333333333333333) - + 2.904737509655563 * f[..., 6] * (x**2 * y - 0.3333333333333333 * y) - + 1.5 * f[..., 3] * x * y - + 0.8660254037844386 * f[..., 2] * y - + 1.677050983124842 * f[..., 4] * (x**2 - 0.3333333333333333) - + 0.8660254037844386 * f[..., 1] * x - + 0.5 * f[..., 0] - ) - - -# end - - -def _expand_2d3p(f, x, y): - return ( - 5.7282196186948 * f[..., 11] * (x * y**3 - 0.6 * x * y) - + 3.307189138830738 * f[..., 9] * (y**3 - 0.6 * y) - + 2.904737509655563 * f[..., 7] * (x * y**2 - 0.3333333333333333 * x) - + 1.677050983124842 * f[..., 5] * (y**2 - 0.3333333333333333) - + 5.7282196186948 * f[..., 10] * (x**3 * y - 0.6 * x * y) - + 2.904737509655563 * f[..., 6] * (x**2 * y - 0.3333333333333333 * y) - + 1.5 * f[..., 3] * x * y - + 0.8660254037844386 * f[..., 2] * y - + 3.307189138830738 * f[..., 8] * (x**3 - 0.6 * x) - + 1.677050983124842 * f[..., 4] * (x**2 - 0.3333333333333333) - + 0.8660254037844386 * f[..., 1] * x - + 0.5 * f[..., 0] - ) - - -# end - - -def _expand_2d4p(f, x, y): - return ( - 11.36658342467076 - * f[..., 16] - * (x * y**4 - 0.8571428571428571 * (x * y**2 - 0.3333333333333333 * x) - 0.2 * x) - + 6.5625 - * f[..., 14] - * (y**4 - 0.8571428571428571 * (y**2 - 0.3333333333333333) - 0.2) - + 5.7282196186948 * f[..., 12] * (x * y**3 - 0.6 * x * y) - + 3.307189138830738 * f[..., 9] * (y**3 - 0.6 * y) - + 5.625 - * f[..., 10] - * ( - x**2 * y**2 - - 0.3333333333333333 * (y**2 - 0.3333333333333333) - - 0.3333333333333333 * (x**2 - 0.3333333333333333) - - 0.1111111111111111 - ) - + 2.904737509655563 * f[..., 7] * (x * y**2 - 0.3333333333333333 * x) - + 1.677050983124842 * f[..., 5] * (y**2 - 0.3333333333333333) - + 11.36658342467076 - * f[..., 15] - * (-0.8571428571428571 * (x**2 * y - 0.3333333333333333 * y) + x**4 * y - 0.2 * y) - + 5.7282196186948 * f[..., 11] * (x**3 * y - 0.6 * x * y) - + 2.904737509655563 * f[..., 6] * (x**2 * y - 0.3333333333333333 * y) - + 1.5 * f[..., 3] * x * y - + 0.8660254037844386 * f[..., 2] * y - + 6.5625 - * f[..., 13] - * (x**4 - 0.8571428571428571 * (x**2 - 0.3333333333333333) - 0.2) - + 3.307189138830738 * f[..., 8] * (x**3 - 0.6 * x) - + 1.677050983124842 * f[..., 4] * (x**2 - 0.3333333333333333) - + 0.8660254037844386 * f[..., 1] * x - + 0.5 * f[..., 0] - ) # end - - -expand_2d = [_expand_2d1p, _expand_2d2p, _expand_2d3p, _expand_2d4p] diff --git a/src/postgkyl/modalDG/kernels/expand3d.py b/src/postgkyl/modalDG/kernels/expand3d.py deleted file mode 100755 index e5433a5c..00000000 --- a/src/postgkyl/modalDG/kernels/expand3d.py +++ /dev/null @@ -1,190 +0,0 @@ -def _expand_3d1p(f, x, y, z): - return ( - 1.837117307087383 * f[..., 7] * x * y * z - + 1.060660171779821 * f[..., 6] * y * z - + 1.060660171779821 * f[..., 5] * x * z - + 0.6123724356957944 * f[..., 3] * z - + 1.060660171779821 * f[..., 4] * x * y - + 0.6123724356957944 * f[..., 2] * y - + 0.6123724356957944 * f[..., 1] * x - + 0.3535533905932737 * f[..., 0] - ) - - -# end - - -def _expand_3d2p(f, x, y, z): - return ( - 3.557562367689425 * f[..., 19] * (x * y * z**2 - 0.3333333333333333 * x * y) - + 2.053959590644372 * f[..., 16] * (y * z**2 - 0.3333333333333333 * y) - + 2.053959590644372 * f[..., 15] * (x * z**2 - 0.3333333333333333 * x) - + 1.185854122563142 * f[..., 9] * (z**2 - 0.3333333333333333) - + 3.557562367689425 * f[..., 18] * (x * y**2 * z - 0.3333333333333333 * x * z) - + 2.053959590644372 * f[..., 14] * (y**2 * z - 0.3333333333333333 * z) - + 3.557562367689425 * f[..., 17] * (x**2 * y * z - 0.3333333333333333 * y * z) - + 2.053959590644372 * f[..., 13] * (x**2 * z - 0.3333333333333333 * z) - + 1.837117307087383 * f[..., 10] * x * y * z - + 1.060660171779821 * f[..., 6] * y * z - + 1.060660171779821 * f[..., 5] * x * z - + 0.6123724356957944 * f[..., 3] * z - + 2.053959590644372 * f[..., 12] * (x * y**2 - 0.3333333333333333 * x) - + 1.185854122563142 * f[..., 8] * (y**2 - 0.3333333333333333) - + 2.053959590644372 * f[..., 11] * (x**2 * y - 0.3333333333333333 * y) - + 1.060660171779821 * f[..., 4] * x * y - + 0.6123724356957944 * f[..., 2] * y - + 1.185854122563142 * f[..., 7] * (x**2 - 0.3333333333333333) - + 0.6123724356957944 * f[..., 1] * x - + 0.3535533905932737 * f[..., 0] - ) - - -# end - - -def _expand_3d3p(f, x, y, z): - return ( - 7.015607600201137 * f[..., 31] * (x * y * z**3 - 0.6 * x * y * z) - + 4.050462936504911 * f[..., 28] * (y * z**3 - 0.6 * y * z) - + 4.050462936504911 * f[..., 27] * (x * z**3 - 0.6 * x * z) - + 2.338535866733713 * f[..., 19] * (z**3 - 0.6 * z) - + 3.557562367689425 * f[..., 22] * (x * y * z**2 - 0.3333333333333333 * x * y) - + 2.053959590644372 * f[..., 16] * (y * z**2 - 0.3333333333333333 * y) - + 2.053959590644372 * f[..., 15] * (x * z**2 - 0.3333333333333333 * x) - + 1.185854122563142 * f[..., 9] * (z**2 - 0.3333333333333333) - + 7.015607600201137 * f[..., 30] * (x * y**3 * z - 0.6 * x * y * z) - + 4.050462936504911 * f[..., 26] * (y**3 * z - 0.6 * y * z) - + 3.557562367689425 * f[..., 21] * (x * y**2 * z - 0.3333333333333333 * x * z) - + 2.053959590644372 * f[..., 14] * (y**2 * z - 0.3333333333333333 * z) - + 7.015607600201137 * f[..., 29] * (x**3 * y * z - 0.6 * x * y * z) - + 3.557562367689425 * f[..., 20] * (x**2 * y * z - 0.3333333333333333 * y * z) - + 4.050462936504911 * f[..., 25] * (x**3 * z - 0.6 * x * z) - + 2.053959590644372 * f[..., 13] * (x**2 * z - 0.3333333333333333 * z) - + 1.837117307087383 * f[..., 10] * x * y * z - + 1.060660171779821 * f[..., 6] * y * z - + 1.060660171779821 * f[..., 5] * x * z - + 0.6123724356957944 * f[..., 3] * z - + 4.050462936504911 * f[..., 24] * (x * y**3 - 0.6 * x * y) - + 2.338535866733713 * f[..., 18] * (y**3 - 0.6 * y) - + 2.053959590644372 * f[..., 12] * (x * y**2 - 0.3333333333333333 * x) - + 1.185854122563142 * f[..., 8] * (y**2 - 0.3333333333333333) - + 4.050462936504911 * f[..., 23] * (x**3 * y - 0.6 * x * y) - + 2.053959590644372 * f[..., 11] * (x**2 * y - 0.3333333333333333 * y) - + 1.060660171779821 * f[..., 4] * x * y - + 0.6123724356957944 * f[..., 2] * y - + 2.338535866733713 * f[..., 17] * (x**3 - 0.6 * x) - + 1.185854122563142 * f[..., 7] * (x**2 - 0.3333333333333333) - + 0.6123724356957944 * f[..., 1] * x - + 0.3535533905932737 * f[..., 0] - ) - - -# end - - -def _expand_3d4p(f, x, y, z): - return ( - 13.92116475461015 - * f[..., 49] - * ( - x * y * z**4 - - 0.8571428571428571 * (x * y * z**2 - 0.3333333333333333 * x * y) - - 0.2 * x * y - ) - + 8.037388218507298 - * f[..., 46] - * (y * z**4 - 0.8571428571428571 * (y * z**2 - 0.3333333333333333 * y) - 0.2 * y) - + 8.037388218507298 - * f[..., 45] - * (x * z**4 - 0.8571428571428571 * (x * z**2 - 0.3333333333333333 * x) - 0.2 * x) - + 4.640388251536716 - * f[..., 34] - * (z**4 - 0.8571428571428571 * (z**2 - 0.3333333333333333) - 0.2) - + 7.015607600201137 * f[..., 40] * (x * y * z**3 - 0.6 * x * y * z) - + 4.050462936504911 * f[..., 31] * (y * z**3 - 0.6 * y * z) - + 4.050462936504911 * f[..., 30] * (x * z**3 - 0.6 * x * z) - + 2.338535866733713 * f[..., 19] * (z**3 - 0.6 * z) - + 6.889189901577683 - * f[..., 36] - * +6.889189901577683 - * f[..., 37] - * +3.977475644174328 - * f[..., 25] - * +3.557562367689425 - * f[..., 22] - * (x * y * z**2 - 0.3333333333333333 * x * y) - + 2.053959590644372 * f[..., 16] * (y * z**2 - 0.3333333333333333 * y) - + 3.977475644174328 - * f[..., 24] - * +2.053959590644372 - * f[..., 15] - * (x * z**2 - 0.3333333333333333 * x) - + 1.185854122563142 * f[..., 9] * (z**2 - 0.3333333333333333) - + 13.92116475461015 - * f[..., 48] - * ( - -0.8571428571428571 * (x * y**2 * z - 0.3333333333333333 * x * z) - + x * y**4 * z - - 0.2 * x * z - ) - + 6.889189901577683 - * f[..., 35] - * +8.037388218507298 - * f[..., 44] - * (-0.8571428571428571 * (y**2 * z - 0.3333333333333333 * z) + y**4 * z - 0.2 * z) - + 13.92116475461015 - * f[..., 47] - * ( - -0.8571428571428571 * (x**2 * y * z - 0.3333333333333333 * y * z) - + x**4 * y * z - - 0.2 * y * z - ) - + 8.037388218507298 - * f[..., 43] - * (-0.8571428571428571 * (x**2 * z - 0.3333333333333333 * z) + x**4 * z - 0.2 * z) - + 7.015607600201137 * f[..., 39] * (x * y**3 * z - 0.6 * x * y * z) - + 4.050462936504911 * f[..., 29] * (y**3 * z - 0.6 * y * z) - + 3.557562367689425 * f[..., 21] * (x * y**2 * z - 0.3333333333333333 * x * z) - + 2.053959590644372 * f[..., 14] * (y**2 * z - 0.3333333333333333 * z) - + 7.015607600201137 * f[..., 38] * (x**3 * y * z - 0.6 * x * y * z) - + 3.557562367689425 * f[..., 20] * (x**2 * y * z - 0.3333333333333333 * y * z) - + 4.050462936504911 * f[..., 28] * (x**3 * z - 0.6 * x * z) - + 2.053959590644372 * f[..., 13] * (x**2 * z - 0.3333333333333333 * z) - + 1.837117307087383 * f[..., 10] * x * y * z - + 1.060660171779821 * f[..., 6] * y * z - + 1.060660171779821 * f[..., 5] * x * z - + 0.6123724356957944 * f[..., 3] * z - + 8.037388218507298 - * f[..., 42] - * (x * y**4 - 0.8571428571428571 * (x * y**2 - 0.3333333333333333 * x) - 0.2 * x) - + 4.640388251536716 - * f[..., 33] - * (y**4 - 0.8571428571428571 * (y**2 - 0.3333333333333333) - 0.2) - + 4.050462936504911 * f[..., 27] * (x * y**3 - 0.6 * x * y) - + 2.338535866733713 * f[..., 18] * (y**3 - 0.6 * y) - + 3.977475644174328 - * f[..., 23] - * +2.053959590644372 - * f[..., 12] - * (x * y**2 - 0.3333333333333333 * x) - + 1.185854122563142 * f[..., 8] * (y**2 - 0.3333333333333333) - + 8.037388218507298 - * f[..., 41] - * (-0.8571428571428571 * (x**2 * y - 0.3333333333333333 * y) + x**4 * y - 0.2 * y) - + 4.050462936504911 * f[..., 26] * (x**3 * y - 0.6 * x * y) - + 2.053959590644372 * f[..., 11] * (x**2 * y - 0.3333333333333333 * y) - + 1.060660171779821 * f[..., 4] * x * y - + 0.6123724356957944 * f[..., 2] * y - + 4.640388251536716 - * f[..., 32] - * (x**4 - 0.8571428571428571 * (x**2 - 0.3333333333333333) - 0.2) - + 2.338535866733713 * f[..., 17] * (x**3 - 0.6 * x) - + 1.185854122563142 * f[..., 7] * (x**2 - 0.3333333333333333) - + 0.6123724356957944 * f[..., 1] * x - + 0.3535533905932737 * f[..., 0] - ) - - -# end - -expand_3d = [_expand_3d1p, _expand_3d2p, _expand_3d3p, _expand_3d4p] diff --git a/src/postgkyl/modalDG/kernels/expand4d.py b/src/postgkyl/modalDG/kernels/expand4d.py deleted file mode 100755 index 81963910..00000000 --- a/src/postgkyl/modalDG/kernels/expand4d.py +++ /dev/null @@ -1,522 +0,0 @@ -def _expand_4d1p(f, x, y, z, vx): - return ( - 2.25 * f[..., 15] * vx * x * y * z - + 1.299038105676658 * f[..., 11] * x * y * z - + 1.299038105676658 * f[..., 14] * vx * y * z - + 0.75 * f[..., 7] * y * z - + 1.299038105676658 * f[..., 13] * vx * x * z - + 0.75 * f[..., 6] * x * z - + 0.75 * f[..., 10] * vx * z - + 0.4330127018922193 * f[..., 3] * z - + 1.299038105676658 * f[..., 12] * vx * x * y - + 0.75 * f[..., 5] * x * y - + 0.75 * f[..., 9] * vx * y - + 0.4330127018922193 * f[..., 2] * y - + 0.75 * f[..., 8] * vx * x - + 0.4330127018922193 * f[..., 1] * x - + 0.4330127018922193 * f[..., 4] * vx - + 0.25 * f[..., 0] - ) - - -# end - - -def _expand_4d2p(f, x, y, z, vx): - return ( - 4.357106264483344 - * f[..., 46] - * (vx * x * y * z**2 - 0.3333333333333333 * vx * x * y) - + 2.515576474687264 * f[..., 34] * (x * y * z**2 - 0.3333333333333333 * x * y) - + 2.515576474687264 * f[..., 40] * (vx * y * z**2 - 0.3333333333333333 * vx * y) - + 1.452368754827781 * f[..., 24] * (y * z**2 - 0.3333333333333333 * y) - + 2.515576474687264 * f[..., 39] * (vx * x * z**2 - 0.3333333333333333 * vx * x) - + 1.452368754827781 * f[..., 23] * (x * z**2 - 0.3333333333333333 * x) - + 1.452368754827781 * f[..., 27] * (vx * z**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 13] * (z**2 - 0.3333333333333333) - + 4.357106264483344 - * f[..., 45] - * (vx * x * y**2 * z - 0.3333333333333333 * vx * x * z) - + 2.515576474687264 * f[..., 33] * (x * y**2 * z - 0.3333333333333333 * x * z) - + 2.515576474687264 * f[..., 38] * (vx * y**2 * z - 0.3333333333333333 * vx * z) - + 1.452368754827781 * f[..., 22] * (y**2 * z - 0.3333333333333333 * z) - + 4.357106264483344 - * f[..., 44] - * (vx * x**2 * y * z - 0.3333333333333333 * vx * y * z) - + 2.515576474687264 * f[..., 32] * (x**2 * y * z - 0.3333333333333333 * y * z) - + 4.357106264483344 - * f[..., 47] - * (vx**2 * x * y * z - 0.3333333333333333 * x * y * z) - + 2.515576474687264 * f[..., 43] * (vx**2 * y * z - 0.3333333333333333 * y * z) - + 2.515576474687264 * f[..., 37] * (vx * x**2 * z - 0.3333333333333333 * vx * z) - + 1.452368754827781 * f[..., 21] * (x**2 * z - 0.3333333333333333 * z) - + 2.515576474687264 * f[..., 42] * (vx**2 * x * z - 0.3333333333333333 * x * z) - + 1.452368754827781 * f[..., 30] * (vx**2 * z - 0.3333333333333333 * z) - + 2.25 * f[..., 31] * vx * x * y * z - + 1.299038105676658 * f[..., 15] * x * y * z - + 1.299038105676658 * f[..., 18] * vx * y * z - + 0.75 * f[..., 7] * y * z - + 1.299038105676658 * f[..., 17] * vx * x * z - + 0.75 * f[..., 6] * x * z - + 0.75 * f[..., 10] * vx * z - + 0.4330127018922193 * f[..., 3] * z - + 2.515576474687264 * f[..., 36] * (vx * x * y**2 - 0.3333333333333333 * vx * x) - + 1.452368754827781 * f[..., 20] * (x * y**2 - 0.3333333333333333 * x) - + 1.452368754827781 * f[..., 26] * (vx * y**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 12] * (y**2 - 0.3333333333333333) - + 2.515576474687264 * f[..., 35] * (vx * x**2 * y - 0.3333333333333333 * vx * y) - + 1.452368754827781 * f[..., 19] * (x**2 * y - 0.3333333333333333 * y) - + 2.515576474687264 * f[..., 41] * (vx**2 * x * y - 0.3333333333333333 * x * y) - + 1.452368754827781 * f[..., 29] * (vx**2 * y - 0.3333333333333333 * y) - + 1.299038105676658 * f[..., 16] * vx * x * y - + 0.75 * f[..., 5] * x * y - + 0.75 * f[..., 9] * vx * y - + 0.4330127018922193 * f[..., 2] * y - + 1.452368754827781 * f[..., 25] * (vx * x**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 11] * (x**2 - 0.3333333333333333) - + 1.452368754827781 * f[..., 28] * (vx**2 * x - 0.3333333333333333 * x) - + 0.75 * f[..., 8] * vx * x - + 0.4330127018922193 * f[..., 1] * x - + 0.8385254915624212 * f[..., 14] * (vx**2 - 0.3333333333333333) - + 0.4330127018922193 * f[..., 4] * vx - + 0.25 * f[..., 0] - ) - - -# end - - -def _expand_4d3p(f, x, y, z, vx): - return ( - 8.5923294280422 * f[..., 78] * (vx * x * y * z**3 - 0.6 * vx * x * y * z) - + 4.960783708246107 * f[..., 66] * (x * y * z**3 - 0.6 * x * y * z) - + 4.960783708246107 * f[..., 72] * (vx * y * z**3 - 0.6 * vx * y * z) - + 2.8641098093474 * f[..., 53] * (y * z**3 - 0.6 * y * z) - + 4.960783708246107 * f[..., 71] * (vx * x * z**3 - 0.6 * vx * x * z) - + 2.8641098093474 * f[..., 52] * (x * z**3 - 0.6 * x * z) - + 2.8641098093474 * f[..., 56] * (vx * z**3 - 0.6 * vx * z) - + 1.653594569415369 * f[..., 33] * (z**3 - 0.6 * z) - + 4.357106264483344 - * f[..., 62] - * (vx * x * y * z**2 - 0.3333333333333333 * vx * x * y) - + 2.515576474687264 * f[..., 38] * (x * y * z**2 - 0.3333333333333333 * x * y) - + 2.515576474687264 * f[..., 44] * (vx * y * z**2 - 0.3333333333333333 * vx * y) - + 1.452368754827781 * f[..., 24] * (y * z**2 - 0.3333333333333333 * y) - + 2.515576474687264 * f[..., 43] * (vx * x * z**2 - 0.3333333333333333 * vx * x) - + 1.452368754827781 * f[..., 23] * (x * z**2 - 0.3333333333333333 * x) - + 1.452368754827781 * f[..., 27] * (vx * z**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 13] * (z**2 - 0.3333333333333333) - + 8.5923294280422 * f[..., 77] * (vx * x * y**3 * z - 0.6 * vx * x * y * z) - + 4.960783708246107 * f[..., 65] * (x * y**3 * z - 0.6 * x * y * z) - + 4.960783708246107 * f[..., 70] * (vx * y**3 * z - 0.6 * vx * y * z) - + 2.8641098093474 * f[..., 51] * (y**3 * z - 0.6 * y * z) - + 4.357106264483344 - * f[..., 61] - * (vx * x * y**2 * z - 0.3333333333333333 * vx * x * z) - + 2.515576474687264 * f[..., 37] * (x * y**2 * z - 0.3333333333333333 * x * z) - + 2.515576474687264 * f[..., 42] * (vx * y**2 * z - 0.3333333333333333 * vx * z) - + 1.452368754827781 * f[..., 22] * (y**2 * z - 0.3333333333333333 * z) - + 8.5923294280422 * f[..., 76] * (vx * x**3 * y * z - 0.6 * vx * x * y * z) - + 4.960783708246107 * f[..., 64] * (x**3 * y * z - 0.6 * x * y * z) - + 4.357106264483344 - * f[..., 60] - * (vx * x**2 * y * z - 0.3333333333333333 * vx * y * z) - + 2.515576474687264 * f[..., 36] * (x**2 * y * z - 0.3333333333333333 * y * z) - + 8.5923294280422 * f[..., 79] * (vx**3 * x * y * z - 0.6 * vx * x * y * z) - + 4.357106264483344 - * f[..., 63] - * (vx**2 * x * y * z - 0.3333333333333333 * x * y * z) - + 4.960783708246107 * f[..., 75] * (vx**3 * y * z - 0.6 * vx * y * z) - + 2.515576474687264 * f[..., 47] * (vx**2 * y * z - 0.3333333333333333 * y * z) - + 4.960783708246107 * f[..., 69] * (vx * x**3 * z - 0.6 * vx * x * z) - + 2.8641098093474 * f[..., 50] * (x**3 * z - 0.6 * x * z) - + 2.515576474687264 * f[..., 41] * (vx * x**2 * z - 0.3333333333333333 * vx * z) - + 1.452368754827781 * f[..., 21] * (x**2 * z - 0.3333333333333333 * z) - + 4.960783708246107 * f[..., 74] * (vx**3 * x * z - 0.6 * vx * x * z) - + 2.515576474687264 * f[..., 46] * (vx**2 * x * z - 0.3333333333333333 * x * z) - + 2.8641098093474 * f[..., 59] * (vx**3 * z - 0.6 * vx * z) - + 1.452368754827781 * f[..., 30] * (vx**2 * z - 0.3333333333333333 * z) - + 2.25 * f[..., 35] * vx * x * y * z - + 1.299038105676658 * f[..., 15] * x * y * z - + 1.299038105676658 * f[..., 18] * vx * y * z - + 0.75 * f[..., 7] * y * z - + 1.299038105676658 * f[..., 17] * vx * x * z - + 0.75 * f[..., 6] * x * z - + 0.75 * f[..., 10] * vx * z - + 0.4330127018922193 * f[..., 3] * z - + 4.960783708246107 * f[..., 68] * (vx * x * y**3 - 0.6 * vx * x * y) - + 2.8641098093474 * f[..., 49] * (x * y**3 - 0.6 * x * y) - + 2.8641098093474 * f[..., 55] * (vx * y**3 - 0.6 * vx * y) - + 1.653594569415369 * f[..., 32] * (y**3 - 0.6 * y) - + 2.515576474687264 * f[..., 40] * (vx * x * y**2 - 0.3333333333333333 * vx * x) - + 1.452368754827781 * f[..., 20] * (x * y**2 - 0.3333333333333333 * x) - + 1.452368754827781 * f[..., 26] * (vx * y**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 12] * (y**2 - 0.3333333333333333) - + 4.960783708246107 * f[..., 67] * (vx * x**3 * y - 0.6 * vx * x * y) - + 2.8641098093474 * f[..., 48] * (x**3 * y - 0.6 * x * y) - + 2.515576474687264 * f[..., 39] * (vx * x**2 * y - 0.3333333333333333 * vx * y) - + 1.452368754827781 * f[..., 19] * (x**2 * y - 0.3333333333333333 * y) - + 4.960783708246107 * f[..., 73] * (vx**3 * x * y - 0.6 * vx * x * y) - + 2.515576474687264 * f[..., 45] * (vx**2 * x * y - 0.3333333333333333 * x * y) - + 2.8641098093474 * f[..., 58] * (vx**3 * y - 0.6 * vx * y) - + 1.452368754827781 * f[..., 29] * (vx**2 * y - 0.3333333333333333 * y) - + 1.299038105676658 * f[..., 16] * vx * x * y - + 0.75 * f[..., 5] * x * y - + 0.75 * f[..., 9] * vx * y - + 0.4330127018922193 * f[..., 2] * y - + 2.8641098093474 * f[..., 54] * (vx * x**3 - 0.6 * vx * x) - + 1.653594569415369 * f[..., 31] * (x**3 - 0.6 * x) - + 1.452368754827781 * f[..., 25] * (vx * x**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 11] * (x**2 - 0.3333333333333333) - + 2.8641098093474 * f[..., 57] * (vx**3 * x - 0.6 * vx * x) - + 1.452368754827781 * f[..., 28] * (vx**2 * x - 0.3333333333333333 * x) - + 0.75 * f[..., 8] * vx * x - + 0.4330127018922193 * f[..., 1] * x - + 1.653594569415369 * f[..., 34] * (vx**3 - 0.6 * vx) - + 0.8385254915624212 * f[..., 14] * (vx**2 - 0.3333333333333333) - + 0.4330127018922193 * f[..., 4] * vx - + 0.25 * f[..., 0] - ) - - -# end - - -def _expand_4d4p(f, x, y, z, vx): - return ( - 17.04987513700613 - * f[..., 134] - * ( - vx * x * y * z**4 - - 0.8571428571428571 * (vx * x * y * z**2 - 0.3333333333333333 * vx * x * y) - - 0.2 * vx * x * y - ) - + 9.84375 - * f[..., 122] - * ( - x * y * z**4 - - 0.8571428571428571 * (x * y * z**2 - 0.3333333333333333 * x * y) - - 0.2 * x * y - ) - + 9.84375 - * f[..., 128] - * ( - vx * y * z**4 - - 0.8571428571428571 * (vx * y * z**2 - 0.3333333333333333 * vx * y) - - 0.2 * vx * y - ) - + 5.683291712335378 - * f[..., 103] - * (y * z**4 - 0.8571428571428571 * (y * z**2 - 0.3333333333333333 * y) - 0.2 * y) - + 9.84375 - * f[..., 127] - * ( - vx * x * z**4 - - 0.8571428571428571 * (vx * x * z**2 - 0.3333333333333333 * vx * x) - - 0.2 * vx * x - ) - + 5.683291712335378 - * f[..., 102] - * (x * z**4 - 0.8571428571428571 * (x * z**2 - 0.3333333333333333 * x) - 0.2 * x) - + 5.683291712335378 - * f[..., 106] - * ( - vx * z**4 - - 0.8571428571428571 * (vx * z**2 - 0.3333333333333333 * vx) - - 0.2 * vx - ) - + 3.28125 - * f[..., 68] - * (z**4 - 0.8571428571428571 * (z**2 - 0.3333333333333333) - 0.2) - + 8.5923294280422 * f[..., 118] * (vx * x * y * z**3 - 0.6 * vx * x * y * z) - + 4.960783708246107 * f[..., 88] * (x * y * z**3 - 0.6 * x * y * z) - + 4.960783708246107 * f[..., 94] * (vx * y * z**3 - 0.6 * vx * y * z) - + 2.8641098093474 * f[..., 59] * (y * z**3 - 0.6 * y * z) - + 4.960783708246107 * f[..., 93] * (vx * x * z**3 - 0.6 * vx * x * z) - + 2.8641098093474 * f[..., 58] * (x * z**3 - 0.6 * x * z) - + 2.8641098093474 * f[..., 62] * (vx * z**3 - 0.6 * vx * z) - + 1.653594569415369 * f[..., 33] * (z**3 - 0.6 * z) - + 8.4375 - * f[..., 115] - * +8.4375 - * f[..., 111] - * +4.871392896287466 - * f[..., 75] - * +4.871392896287466 - * f[..., 85] - * +8.4375 - * f[..., 112] - * +4.871392896287466 - * f[..., 76] - * +4.871392896287466 - * f[..., 84] - * +4.871392896287466 - * f[..., 79] - * +4.871392896287466 - * f[..., 78] - * +2.8125 - * f[..., 50] - * +4.357106264483344 - * f[..., 72] - * (vx * x * y * z**2 - 0.3333333333333333 * vx * x * y) - + 2.515576474687264 * f[..., 38] * (x * y * z**2 - 0.3333333333333333 * x * y) - + 2.515576474687264 * f[..., 44] * (vx * y * z**2 - 0.3333333333333333 * vx * y) - + 1.452368754827781 * f[..., 24] * (y * z**2 - 0.3333333333333333 * y) - + 2.8125 - * f[..., 49] - * +2.515576474687264 - * f[..., 43] - * (vx * x * z**2 - 0.3333333333333333 * vx * x) - + 1.452368754827781 * f[..., 23] * (x * z**2 - 0.3333333333333333 * x) - + 2.8125 - * f[..., 53] - * +1.452368754827781 - * f[..., 27] - * (vx * z**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 13] * (z**2 - 0.3333333333333333) - + 17.04987513700613 - * f[..., 133] - * ( - -0.8571428571428571 * (vx * x * y**2 * z - 0.3333333333333333 * vx * x * z) - + vx * x * y**4 * z - - 0.2 * vx * x * z - ) - + 8.4375 - * f[..., 114] - * +9.84375 - * f[..., 121] - * ( - -0.8571428571428571 * (x * y**2 * z - 0.3333333333333333 * x * z) - + x * y**4 * z - - 0.2 * x * z - ) - + 8.4375 - * f[..., 110] - * +9.84375 - * f[..., 126] - * ( - -0.8571428571428571 * (vx * y**2 * z - 0.3333333333333333 * vx * z) - + vx * y**4 * z - - 0.2 * vx * z - ) - + 4.871392896287466 - * f[..., 74] - * +4.871392896287466 - * f[..., 83] - * +5.683291712335378 - * f[..., 101] - * (-0.8571428571428571 * (y**2 * z - 0.3333333333333333 * z) + y**4 * z - 0.2 * z) - + 17.04987513700613 - * f[..., 132] - * ( - -0.8571428571428571 * (vx * x**2 * y * z - 0.3333333333333333 * vx * y * z) - + vx * x**4 * y * z - - 0.2 * vx * y * z - ) - + 8.4375 - * f[..., 113] - * +9.84375 - * f[..., 120] - * ( - -0.8571428571428571 * (x**2 * y * z - 0.3333333333333333 * y * z) - + x**4 * y * z - - 0.2 * y * z - ) - + 17.04987513700613 - * f[..., 135] - * ( - -0.8571428571428571 * (vx**2 * x * y * z - 0.3333333333333333 * x * y * z) - + vx**4 * x * y * z - - 0.2 * x * y * z - ) - + 9.84375 - * f[..., 131] - * ( - -0.8571428571428571 * (vx**2 * y * z - 0.3333333333333333 * y * z) - + vx**4 * y * z - - 0.2 * y * z - ) - + 9.84375 - * f[..., 125] - * ( - -0.8571428571428571 * (vx * x**2 * z - 0.3333333333333333 * vx * z) - + vx * x**4 * z - - 0.2 * vx * z - ) - + 4.871392896287466 - * f[..., 82] - * +5.683291712335378 - * f[..., 100] - * (-0.8571428571428571 * (x**2 * z - 0.3333333333333333 * z) + x**4 * z - 0.2 * z) - + 9.84375 - * f[..., 130] - * ( - -0.8571428571428571 * (vx**2 * x * z - 0.3333333333333333 * x * z) - + vx**4 * x * z - - 0.2 * x * z - ) - + 5.683291712335378 - * f[..., 109] - * ( - -0.8571428571428571 * (vx**2 * z - 0.3333333333333333 * z) - + vx**4 * z - - 0.2 * z - ) - + 8.5923294280422 * f[..., 117] * (vx * x * y**3 * z - 0.6 * vx * x * y * z) - + 4.960783708246107 * f[..., 87] * (x * y**3 * z - 0.6 * x * y * z) - + 4.960783708246107 * f[..., 92] * (vx * y**3 * z - 0.6 * vx * y * z) - + 2.8641098093474 * f[..., 57] * (y**3 * z - 0.6 * y * z) - + 4.357106264483344 - * f[..., 71] - * (vx * x * y**2 * z - 0.3333333333333333 * vx * x * z) - + 2.515576474687264 * f[..., 37] * (x * y**2 * z - 0.3333333333333333 * x * z) - + 2.515576474687264 * f[..., 42] * (vx * y**2 * z - 0.3333333333333333 * vx * z) - + 1.452368754827781 * f[..., 22] * (y**2 * z - 0.3333333333333333 * z) - + 8.5923294280422 * f[..., 116] * (vx * x**3 * y * z - 0.6 * vx * x * y * z) - + 4.960783708246107 * f[..., 86] * (x**3 * y * z - 0.6 * x * y * z) - + 4.357106264483344 - * f[..., 70] - * (vx * x**2 * y * z - 0.3333333333333333 * vx * y * z) - + 2.515576474687264 * f[..., 36] * (x**2 * y * z - 0.3333333333333333 * y * z) - + 8.5923294280422 * f[..., 119] * (vx**3 * x * y * z - 0.6 * vx * x * y * z) - + 4.357106264483344 - * f[..., 73] - * (vx**2 * x * y * z - 0.3333333333333333 * x * y * z) - + 4.960783708246107 * f[..., 97] * (vx**3 * y * z - 0.6 * vx * y * z) - + 2.515576474687264 * f[..., 47] * (vx**2 * y * z - 0.3333333333333333 * y * z) - + 4.960783708246107 * f[..., 91] * (vx * x**3 * z - 0.6 * vx * x * z) - + 2.8641098093474 * f[..., 56] * (x**3 * z - 0.6 * x * z) - + 2.515576474687264 * f[..., 41] * (vx * x**2 * z - 0.3333333333333333 * vx * z) - + 1.452368754827781 * f[..., 21] * (x**2 * z - 0.3333333333333333 * z) - + 4.960783708246107 * f[..., 96] * (vx**3 * x * z - 0.6 * vx * x * z) - + 2.515576474687264 * f[..., 46] * (vx**2 * x * z - 0.3333333333333333 * x * z) - + 2.8641098093474 * f[..., 65] * (vx**3 * z - 0.6 * vx * z) - + 1.452368754827781 * f[..., 30] * (vx**2 * z - 0.3333333333333333 * z) - + 2.25 * f[..., 35] * vx * x * y * z - + 1.299038105676658 * f[..., 15] * x * y * z - + 1.299038105676658 * f[..., 18] * vx * y * z - + 0.75 * f[..., 7] * y * z - + 1.299038105676658 * f[..., 17] * vx * x * z - + 0.75 * f[..., 6] * x * z - + 0.75 * f[..., 10] * vx * z - + 0.4330127018922193 * f[..., 3] * z - + 9.84375 - * f[..., 124] - * ( - vx * x * y**4 - - 0.8571428571428571 * (vx * x * y**2 - 0.3333333333333333 * vx * x) - - 0.2 * vx * x - ) - + 5.683291712335378 - * f[..., 99] - * (x * y**4 - 0.8571428571428571 * (x * y**2 - 0.3333333333333333 * x) - 0.2 * x) - + 5.683291712335378 - * f[..., 105] - * ( - vx * y**4 - - 0.8571428571428571 * (vx * y**2 - 0.3333333333333333 * vx) - - 0.2 * vx - ) - + 3.28125 - * f[..., 67] - * (y**4 - 0.8571428571428571 * (y**2 - 0.3333333333333333) - 0.2) - + 4.960783708246107 * f[..., 90] * (vx * x * y**3 - 0.6 * vx * x * y) - + 2.8641098093474 * f[..., 55] * (x * y**3 - 0.6 * x * y) - + 2.8641098093474 * f[..., 61] * (vx * y**3 - 0.6 * vx * y) - + 1.653594569415369 * f[..., 32] * (y**3 - 0.6 * y) - + 4.871392896287466 - * f[..., 81] - * +4.871392896287466 - * f[..., 77] - * +2.8125 - * f[..., 48] - * +2.515576474687264 - * f[..., 40] - * (vx * x * y**2 - 0.3333333333333333 * vx * x) - + 1.452368754827781 * f[..., 20] * (x * y**2 - 0.3333333333333333 * x) - + 2.8125 - * f[..., 52] - * +1.452368754827781 - * f[..., 26] - * (vx * y**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 12] * (y**2 - 0.3333333333333333) - + 9.84375 - * f[..., 123] - * ( - -0.8571428571428571 * (vx * x**2 * y - 0.3333333333333333 * vx * y) - + vx * x**4 * y - - 0.2 * vx * y - ) - + 4.871392896287466 - * f[..., 80] - * +5.683291712335378 - * f[..., 98] - * (-0.8571428571428571 * (x**2 * y - 0.3333333333333333 * y) + x**4 * y - 0.2 * y) - + 9.84375 - * f[..., 129] - * ( - -0.8571428571428571 * (vx**2 * x * y - 0.3333333333333333 * x * y) - + vx**4 * x * y - - 0.2 * x * y - ) - + 5.683291712335378 - * f[..., 108] - * ( - -0.8571428571428571 * (vx**2 * y - 0.3333333333333333 * y) - + vx**4 * y - - 0.2 * y - ) - + 4.960783708246107 * f[..., 89] * (vx * x**3 * y - 0.6 * vx * x * y) - + 2.8641098093474 * f[..., 54] * (x**3 * y - 0.6 * x * y) - + 2.515576474687264 * f[..., 39] * (vx * x**2 * y - 0.3333333333333333 * vx * y) - + 1.452368754827781 * f[..., 19] * (x**2 * y - 0.3333333333333333 * y) - + 4.960783708246107 * f[..., 95] * (vx**3 * x * y - 0.6 * vx * x * y) - + 2.515576474687264 * f[..., 45] * (vx**2 * x * y - 0.3333333333333333 * x * y) - + 2.8641098093474 * f[..., 64] * (vx**3 * y - 0.6 * vx * y) - + 1.452368754827781 * f[..., 29] * (vx**2 * y - 0.3333333333333333 * y) - + 1.299038105676658 * f[..., 16] * vx * x * y - + 0.75 * f[..., 5] * x * y - + 0.75 * f[..., 9] * vx * y - + 0.4330127018922193 * f[..., 2] * y - + 5.683291712335378 - * f[..., 104] - * ( - vx * x**4 - - 0.8571428571428571 * (vx * x**2 - 0.3333333333333333 * vx) - - 0.2 * vx - ) - + 3.28125 - * f[..., 66] - * (x**4 - 0.8571428571428571 * (x**2 - 0.3333333333333333) - 0.2) - + 2.8641098093474 * f[..., 60] * (vx * x**3 - 0.6 * vx * x) - + 1.653594569415369 * f[..., 31] * (x**3 - 0.6 * x) - + 2.8125 - * f[..., 51] - * +1.452368754827781 - * f[..., 25] - * (vx * x**2 - 0.3333333333333333 * vx) - + 0.8385254915624212 * f[..., 11] * (x**2 - 0.3333333333333333) - + 5.683291712335378 - * f[..., 107] - * ( - -0.8571428571428571 * (vx**2 * x - 0.3333333333333333 * x) - + vx**4 * x - - 0.2 * x - ) - + 2.8641098093474 * f[..., 63] * (vx**3 * x - 0.6 * vx * x) - + 1.452368754827781 * f[..., 28] * (vx**2 * x - 0.3333333333333333 * x) - + 0.75 * f[..., 8] * vx * x - + 0.4330127018922193 * f[..., 1] * x - + 3.28125 - * f[..., 69] - * (vx**4 - 0.8571428571428571 * (vx**2 - 0.3333333333333333) - 0.2) - + 1.653594569415369 * f[..., 34] * (vx**3 - 0.6 * vx) - + 0.8385254915624212 * f[..., 14] * (vx**2 - 0.3333333333333333) - + 0.4330127018922193 * f[..., 4] * vx - + 0.25 * f[..., 0] - ) - - -# end - -expand_4d = [_expand_4d1p, _expand_4d2p, _expand_4d3p, _expand_4d4p] diff --git a/src/postgkyl/modalDG/kernels/expand5d.py b/src/postgkyl/modalDG/kernels/expand5d.py deleted file mode 100755 index acfe0e8a..00000000 --- a/src/postgkyl/modalDG/kernels/expand5d.py +++ /dev/null @@ -1,1412 +0,0 @@ -def _expand_5d1p(f, x, y, z, vx, vy): - return ( - 2.755675960631073 * f[..., 31] * vx * vy * x * y * z - + 1.590990257669731 * f[..., 27] * vy * x * y * z - + 1.590990257669731 * f[..., 26] * vx * x * y * z - + 0.9185586535436913 * f[..., 16] * x * y * z - + 1.590990257669731 * f[..., 30] * vx * vy * y * z - + 0.9185586535436913 * f[..., 22] * vy * y * z - + 0.9185586535436913 * f[..., 19] * vx * y * z - + 0.5303300858899105 * f[..., 8] * y * z - + 1.590990257669731 * f[..., 29] * vx * vy * x * z - + 0.9185586535436913 * f[..., 21] * vy * x * z - + 0.9185586535436913 * f[..., 18] * vx * x * z - + 0.5303300858899105 * f[..., 7] * x * z - + 0.9185586535436913 * f[..., 25] * vx * vy * z - + 0.5303300858899105 * f[..., 14] * vy * z - + 0.5303300858899105 * f[..., 11] * vx * z - + 0.3061862178478971 * f[..., 3] * z - + 1.590990257669731 * f[..., 28] * vx * vy * x * y - + 0.9185586535436913 * f[..., 20] * vy * x * y - + 0.9185586535436913 * f[..., 17] * vx * x * y - + 0.5303300858899105 * f[..., 6] * x * y - + 0.9185586535436913 * f[..., 24] * vx * vy * y - + 0.5303300858899105 * f[..., 13] * vy * y - + 0.5303300858899105 * f[..., 10] * vx * y - + 0.3061862178478971 * f[..., 2] * y - + 0.9185586535436913 * f[..., 23] * vx * vy * x - + 0.5303300858899105 * f[..., 12] * vy * x - + 0.5303300858899105 * f[..., 9] * vx * x - + 0.3061862178478971 * f[..., 1] * x - + 0.5303300858899105 * f[..., 15] * vx * vy - + 0.3061862178478971 * f[..., 5] * vy - + 0.3061862178478971 * f[..., 4] * vx - + 0.1767766952966368 * f[..., 0] - ) - - -# end - - -def _expand_5d2p(f, x, y, z, vx, vy): - return ( - 5.336343551534138 - * f[..., 109] - * (vx * vy * x * y * z ^ 2 - 0.3333333333333333 * vx * vy * x * y) - + 3.080939385966558 - * f[..., 93] - * (vy * x * y * z ^ 2 - 0.3333333333333333 * vy * x * y) - + 3.080939385966558 - * f[..., 89] - * (vx * x * y * z ^ 2 - 0.3333333333333333 * vx * x * y) - + 1.778781183844713 * f[..., 58] * (x * y * z ^ 2 - 0.3333333333333333 * x * y) - + 3.080939385966558 - * f[..., 99] - * (vx * vy * y * z ^ 2 - 0.3333333333333333 * vx * vy * y) - + 1.778781183844713 * f[..., 73] * (vy * y * z ^ 2 - 0.3333333333333333 * vy * y) - + 1.778781183844713 * f[..., 64] * (vx * y * z ^ 2 - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 36] * (y * z ^ 2 - 0.3333333333333333 * y) - + 3.080939385966558 - * f[..., 98] - * (vx * vy * x * z ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.778781183844713 * f[..., 72] * (vy * x * z ^ 2 - 0.3333333333333333 * vy * x) - + 1.778781183844713 * f[..., 63] * (vx * x * z ^ 2 - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 35] * (x * z ^ 2 - 0.3333333333333333 * x) - + 1.778781183844713 - * f[..., 76] - * (vx * vy * z ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 45] * (vy * z ^ 2 - 0.3333333333333333 * vy) - + 1.026979795322186 * f[..., 39] * (vx * z ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 18] * (z ^ 2 - 0.3333333333333333) - + 5.336343551534138 - * f[..., 108] - * (vx * vy * x * y ^ 2 * z - 0.3333333333333333 * vx * vy * x * z) - + 3.080939385966558 - * f[..., 92] - * (vy * x * y ^ 2 * z - 0.3333333333333333 * vy * x * z) - + 3.080939385966558 - * f[..., 88] - * (vx * x * y ^ 2 * z - 0.3333333333333333 * vx * x * z) - + 1.778781183844713 * f[..., 57] * (x * y ^ 2 * z - 0.3333333333333333 * x * z) - + 3.080939385966558 - * f[..., 97] - * (vx * vy * y ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.778781183844713 * f[..., 71] * (vy * y ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.778781183844713 * f[..., 62] * (vx * y ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 34] * (y ^ 2 * z - 0.3333333333333333 * z) - + 5.336343551534138 - * f[..., 107] - * (vx * vy * x ^ 2 * y * z - 0.3333333333333333 * vx * vy * y * z) - + 3.080939385966558 - * f[..., 91] - * (vy * x ^ 2 * y * z - 0.3333333333333333 * vy * y * z) - + 3.080939385966558 - * f[..., 87] - * (vx * x ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.778781183844713 * f[..., 56] * (x ^ 2 * y * z - 0.3333333333333333 * y * z) - + 5.336343551534138 - * f[..., 111] - * (vx * vy ^ 2 * x * y * z - 0.3333333333333333 * vx * x * y * z) - + 3.080939385966558 - * f[..., 103] - * (vy ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 5.336343551534138 - * f[..., 110] - * (vx ^ 2 * vy * x * y * z - 0.3333333333333333 * vy * x * y * z) - + 3.080939385966558 - * f[..., 90] - * (vx ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 3.080939385966558 - * f[..., 106] - * (vx * vy ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.778781183844713 * f[..., 82] * (vy ^ 2 * y * z - 0.3333333333333333 * y * z) - + 3.080939385966558 - * f[..., 102] - * (vx ^ 2 * vy * y * z - 0.3333333333333333 * vy * y * z) - + 1.778781183844713 * f[..., 67] * (vx ^ 2 * y * z - 0.3333333333333333 * y * z) - + 3.080939385966558 - * f[..., 96] - * (vx * vy * x ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.778781183844713 * f[..., 70] * (vy * x ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.778781183844713 * f[..., 61] * (vx * x ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 33] * (x ^ 2 * z - 0.3333333333333333 * z) - + 3.080939385966558 - * f[..., 105] - * (vx * vy ^ 2 * x * z - 0.3333333333333333 * vx * x * z) - + 1.778781183844713 * f[..., 81] * (vy ^ 2 * x * z - 0.3333333333333333 * x * z) - + 3.080939385966558 - * f[..., 101] - * (vx ^ 2 * vy * x * z - 0.3333333333333333 * vy * x * z) - + 1.778781183844713 * f[..., 66] * (vx ^ 2 * x * z - 0.3333333333333333 * x * z) - + 1.778781183844713 * f[..., 85] * (vx * vy ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 49] * (vy ^ 2 * z - 0.3333333333333333 * z) - + 1.778781183844713 * f[..., 79] * (vx ^ 2 * vy * z - 0.3333333333333333 * vy * z) - + 1.026979795322186 * f[..., 42] * (vx ^ 2 * z - 0.3333333333333333 * z) - + 2.755675960631073 * f[..., 86] * vx * vy * x * y * z - + 1.590990257669731 * f[..., 52] * vy * x * y * z - + 1.590990257669731 * f[..., 51] * vx * x * y * z - + 0.9185586535436913 * f[..., 21] * x * y * z - + 1.590990257669731 * f[..., 55] * vx * vy * y * z - + 0.9185586535436913 * f[..., 27] * vy * y * z - + 0.9185586535436913 * f[..., 24] * vx * y * z - + 0.5303300858899105 * f[..., 8] * y * z - + 1.590990257669731 * f[..., 54] * vx * vy * x * z - + 0.9185586535436913 * f[..., 26] * vy * x * z - + 0.9185586535436913 * f[..., 23] * vx * x * z - + 0.5303300858899105 * f[..., 7] * x * z - + 0.9185586535436913 * f[..., 30] * vx * vy * z - + 0.5303300858899105 * f[..., 14] * vy * z - + 0.5303300858899105 * f[..., 11] * vx * z - + 0.3061862178478971 * f[..., 3] * z - + 3.080939385966558 - * f[..., 95] - * (vx * vy * x * y ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.778781183844713 * f[..., 69] * (vy * x * y ^ 2 - 0.3333333333333333 * vy * x) - + 1.778781183844713 * f[..., 60] * (vx * x * y ^ 2 - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 32] * (x * y ^ 2 - 0.3333333333333333 * x) - + 1.778781183844713 - * f[..., 75] - * (vx * vy * y ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 44] * (vy * y ^ 2 - 0.3333333333333333 * vy) - + 1.026979795322186 * f[..., 38] * (vx * y ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 17] * (y ^ 2 - 0.3333333333333333) - + 3.080939385966558 - * f[..., 94] - * (vx * vy * x ^ 2 * y - 0.3333333333333333 * vx * vy * y) - + 1.778781183844713 * f[..., 68] * (vy * x ^ 2 * y - 0.3333333333333333 * vy * y) - + 1.778781183844713 * f[..., 59] * (vx * x ^ 2 * y - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 31] * (x ^ 2 * y - 0.3333333333333333 * y) - + 3.080939385966558 - * f[..., 104] - * (vx * vy ^ 2 * x * y - 0.3333333333333333 * vx * x * y) - + 1.778781183844713 * f[..., 80] * (vy ^ 2 * x * y - 0.3333333333333333 * x * y) - + 3.080939385966558 - * f[..., 100] - * (vx ^ 2 * vy * x * y - 0.3333333333333333 * vy * x * y) - + 1.778781183844713 * f[..., 65] * (vx ^ 2 * x * y - 0.3333333333333333 * x * y) - + 1.778781183844713 * f[..., 84] * (vx * vy ^ 2 * y - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 48] * (vy ^ 2 * y - 0.3333333333333333 * y) - + 1.778781183844713 * f[..., 78] * (vx ^ 2 * vy * y - 0.3333333333333333 * vy * y) - + 1.026979795322186 * f[..., 41] * (vx ^ 2 * y - 0.3333333333333333 * y) - + 1.590990257669731 * f[..., 53] * vx * vy * x * y - + 0.9185586535436913 * f[..., 25] * vy * x * y - + 0.9185586535436913 * f[..., 22] * vx * x * y - + 0.5303300858899105 * f[..., 6] * x * y - + 0.9185586535436913 * f[..., 29] * vx * vy * y - + 0.5303300858899105 * f[..., 13] * vy * y - + 0.5303300858899105 * f[..., 10] * vx * y - + 0.3061862178478971 * f[..., 2] * y - + 1.778781183844713 - * f[..., 74] - * (vx * vy * x ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 43] * (vy * x ^ 2 - 0.3333333333333333 * vy) - + 1.026979795322186 * f[..., 37] * (vx * x ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 16] * (x ^ 2 - 0.3333333333333333) - + 1.778781183844713 * f[..., 83] * (vx * vy ^ 2 * x - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 47] * (vy ^ 2 * x - 0.3333333333333333 * x) - + 1.778781183844713 * f[..., 77] * (vx ^ 2 * vy * x - 0.3333333333333333 * vy * x) - + 1.026979795322186 * f[..., 40] * (vx ^ 2 * x - 0.3333333333333333 * x) - + 0.9185586535436913 * f[..., 28] * vx * vy * x - + 0.5303300858899105 * f[..., 12] * vy * x - + 0.5303300858899105 * f[..., 9] * vx * x - + 0.3061862178478971 * f[..., 1] * x - + 1.026979795322186 * f[..., 50] * (vx * vy ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 20] * (vy ^ 2 - 0.3333333333333333) - + 1.026979795322186 * f[..., 46] * (vx ^ 2 * vy - 0.3333333333333333 * vy) - + 0.5303300858899105 * f[..., 15] * vx * vy - + 0.3061862178478971 * f[..., 5] * vy - + 0.592927061281571 * f[..., 19] * (vx ^ 2 - 0.3333333333333333) - + 0.3061862178478971 * f[..., 4] * vx - + 0.1767766952966368 * f[..., 0] - ) - - -# end - - -def _expand_5d3p(f, x, y, z, vx, vy): - return ( - 10.52341140030171 - * f[..., 189] - * (vx * vy * x * y * z ^ 3 - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 173] * (vy * x * y * z ^ 3 - 0.6 * vy * x * y * z) - + 6.075694404757366 * f[..., 169] * (vx * x * y * z ^ 3 - 0.6 * vx * x * y * z) - + 3.507803800100568 * f[..., 134] * (x * y * z ^ 3 - 0.6 * x * y * z) - + 6.075694404757366 * f[..., 179] * (vx * vy * y * z ^ 3 - 0.6 * vx * vy * y * z) - + 3.507803800100568 * f[..., 149] * (vy * y * z ^ 3 - 0.6 * vy * y * z) - + 3.507803800100568 * f[..., 140] * (vx * y * z ^ 3 - 0.6 * vx * y * z) - + 2.025231468252455 * f[..., 96] * (y * z ^ 3 - 0.6 * y * z) - + 6.075694404757366 * f[..., 178] * (vx * vy * x * z ^ 3 - 0.6 * vx * vy * x * z) - + 3.507803800100568 * f[..., 148] * (vy * x * z ^ 3 - 0.6 * vy * x * z) - + 3.507803800100568 * f[..., 139] * (vx * x * z ^ 3 - 0.6 * vx * x * z) - + 2.025231468252455 * f[..., 95] * (x * z ^ 3 - 0.6 * x * z) - + 3.507803800100568 * f[..., 152] * (vx * vy * z ^ 3 - 0.6 * vx * vy * z) - + 2.025231468252455 * f[..., 105] * (vy * z ^ 3 - 0.6 * vy * z) - + 2.025231468252455 * f[..., 99] * (vx * z ^ 3 - 0.6 * vx * z) - + 1.169267933366856 * f[..., 53] * (z ^ 3 - 0.6 * z) - + 5.336343551534138 - * f[..., 164] - * (vx * vy * x * y * z ^ 2 - 0.3333333333333333 * vx * vy * x * y) - + 3.080939385966558 - * f[..., 118] - * (vy * x * y * z ^ 2 - 0.3333333333333333 * vy * x * y) - + 3.080939385966558 - * f[..., 114] - * (vx * x * y * z ^ 2 - 0.3333333333333333 * vx * x * y) - + 1.778781183844713 * f[..., 63] * (x * y * z ^ 2 - 0.3333333333333333 * x * y) - + 3.080939385966558 - * f[..., 124] - * (vx * vy * y * z ^ 2 - 0.3333333333333333 * vx * vy * y) - + 1.778781183844713 * f[..., 78] * (vy * y * z ^ 2 - 0.3333333333333333 * vy * y) - + 1.778781183844713 * f[..., 69] * (vx * y * z ^ 2 - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 36] * (y * z ^ 2 - 0.3333333333333333 * y) - + 3.080939385966558 - * f[..., 123] - * (vx * vy * x * z ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.778781183844713 * f[..., 77] * (vy * x * z ^ 2 - 0.3333333333333333 * vy * x) - + 1.778781183844713 * f[..., 68] * (vx * x * z ^ 2 - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 35] * (x * z ^ 2 - 0.3333333333333333 * x) - + 1.778781183844713 - * f[..., 81] - * (vx * vy * z ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 45] * (vy * z ^ 2 - 0.3333333333333333 * vy) - + 1.026979795322186 * f[..., 39] * (vx * z ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 18] * (z ^ 2 - 0.3333333333333333) - + 10.52341140030171 - * f[..., 188] - * (vx * vy * x * y ^ 3 * z - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 172] * (vy * x * y ^ 3 * z - 0.6 * vy * x * y * z) - + 6.075694404757366 * f[..., 168] * (vx * x * y ^ 3 * z - 0.6 * vx * x * y * z) - + 3.507803800100568 * f[..., 133] * (x * y ^ 3 * z - 0.6 * x * y * z) - + 6.075694404757366 * f[..., 177] * (vx * vy * y ^ 3 * z - 0.6 * vx * vy * y * z) - + 3.507803800100568 * f[..., 147] * (vy * y ^ 3 * z - 0.6 * vy * y * z) - + 3.507803800100568 * f[..., 138] * (vx * y ^ 3 * z - 0.6 * vx * y * z) - + 2.025231468252455 * f[..., 94] * (y ^ 3 * z - 0.6 * y * z) - + 5.336343551534138 - * f[..., 163] - * (vx * vy * x * y ^ 2 * z - 0.3333333333333333 * vx * vy * x * z) - + 3.080939385966558 - * f[..., 117] - * (vy * x * y ^ 2 * z - 0.3333333333333333 * vy * x * z) - + 3.080939385966558 - * f[..., 113] - * (vx * x * y ^ 2 * z - 0.3333333333333333 * vx * x * z) - + 1.778781183844713 * f[..., 62] * (x * y ^ 2 * z - 0.3333333333333333 * x * z) - + 3.080939385966558 - * f[..., 122] - * (vx * vy * y ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.778781183844713 * f[..., 76] * (vy * y ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.778781183844713 * f[..., 67] * (vx * y ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 34] * (y ^ 2 * z - 0.3333333333333333 * z) - + 10.52341140030171 - * f[..., 187] - * (vx * vy * x ^ 3 * y * z - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 171] * (vy * x ^ 3 * y * z - 0.6 * vy * x * y * z) - + 6.075694404757366 * f[..., 167] * (vx * x ^ 3 * y * z - 0.6 * vx * x * y * z) - + 3.507803800100568 * f[..., 132] * (x ^ 3 * y * z - 0.6 * x * y * z) - + 5.336343551534138 - * f[..., 162] - * (vx * vy * x ^ 2 * y * z - 0.3333333333333333 * vx * vy * y * z) - + 3.080939385966558 - * f[..., 116] - * (vy * x ^ 2 * y * z - 0.3333333333333333 * vy * y * z) - + 3.080939385966558 - * f[..., 112] - * (vx * x ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.778781183844713 * f[..., 61] * (x ^ 2 * y * z - 0.3333333333333333 * y * z) - + 10.52341140030171 - * f[..., 191] - * (vx * vy ^ 3 * x * y * z - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 183] * (vy ^ 3 * x * y * z - 0.6 * vy * x * y * z) - + 5.336343551534138 - * f[..., 166] - * (vx * vy ^ 2 * x * y * z - 0.3333333333333333 * vx * x * y * z) - + 3.080939385966558 - * f[..., 128] - * (vy ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 10.52341140030171 - * f[..., 190] - * (vx ^ 3 * vy * x * y * z - 0.6 * vx * vy * x * y * z) - + 5.336343551534138 - * f[..., 165] - * (vx ^ 2 * vy * x * y * z - 0.3333333333333333 * vy * x * y * z) - + 6.075694404757366 * f[..., 170] * (vx ^ 3 * x * y * z - 0.6 * vx * x * y * z) - + 3.080939385966558 - * f[..., 115] - * (vx ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 6.075694404757366 * f[..., 186] * (vx * vy ^ 3 * y * z - 0.6 * vx * vy * y * z) - + 3.507803800100568 * f[..., 158] * (vy ^ 3 * y * z - 0.6 * vy * y * z) - + 3.080939385966558 - * f[..., 131] - * (vx * vy ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.778781183844713 * f[..., 87] * (vy ^ 2 * y * z - 0.3333333333333333 * y * z) - + 6.075694404757366 * f[..., 182] * (vx ^ 3 * vy * y * z - 0.6 * vx * vy * y * z) - + 3.080939385966558 - * f[..., 127] - * (vx ^ 2 * vy * y * z - 0.3333333333333333 * vy * y * z) - + 3.507803800100568 * f[..., 143] * (vx ^ 3 * y * z - 0.6 * vx * y * z) - + 1.778781183844713 * f[..., 72] * (vx ^ 2 * y * z - 0.3333333333333333 * y * z) - + 6.075694404757366 * f[..., 176] * (vx * vy * x ^ 3 * z - 0.6 * vx * vy * x * z) - + 3.507803800100568 * f[..., 146] * (vy * x ^ 3 * z - 0.6 * vy * x * z) - + 3.507803800100568 * f[..., 137] * (vx * x ^ 3 * z - 0.6 * vx * x * z) - + 2.025231468252455 * f[..., 93] * (x ^ 3 * z - 0.6 * x * z) - + 3.080939385966558 - * f[..., 121] - * (vx * vy * x ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.778781183844713 * f[..., 75] * (vy * x ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.778781183844713 * f[..., 66] * (vx * x ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 33] * (x ^ 2 * z - 0.3333333333333333 * z) - + 6.075694404757366 * f[..., 185] * (vx * vy ^ 3 * x * z - 0.6 * vx * vy * x * z) - + 3.507803800100568 * f[..., 157] * (vy ^ 3 * x * z - 0.6 * vy * x * z) - + 3.080939385966558 - * f[..., 130] - * (vx * vy ^ 2 * x * z - 0.3333333333333333 * vx * x * z) - + 1.778781183844713 * f[..., 86] * (vy ^ 2 * x * z - 0.3333333333333333 * x * z) - + 6.075694404757366 * f[..., 181] * (vx ^ 3 * vy * x * z - 0.6 * vx * vy * x * z) - + 3.080939385966558 - * f[..., 126] - * (vx ^ 2 * vy * x * z - 0.3333333333333333 * vy * x * z) - + 3.507803800100568 * f[..., 142] * (vx ^ 3 * x * z - 0.6 * vx * x * z) - + 1.778781183844713 * f[..., 71] * (vx ^ 2 * x * z - 0.3333333333333333 * x * z) - + 3.507803800100568 * f[..., 161] * (vx * vy ^ 3 * z - 0.6 * vx * vy * z) - + 2.025231468252455 * f[..., 109] * (vy ^ 3 * z - 0.6 * vy * z) - + 1.778781183844713 * f[..., 90] * (vx * vy ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 49] * (vy ^ 2 * z - 0.3333333333333333 * z) - + 3.507803800100568 * f[..., 155] * (vx ^ 3 * vy * z - 0.6 * vx * vy * z) - + 1.778781183844713 * f[..., 84] * (vx ^ 2 * vy * z - 0.3333333333333333 * vy * z) - + 2.025231468252455 * f[..., 102] * (vx ^ 3 * z - 0.6 * vx * z) - + 1.026979795322186 * f[..., 42] * (vx ^ 2 * z - 0.3333333333333333 * z) - + 2.755675960631073 * f[..., 111] * vx * vy * x * y * z - + 1.590990257669731 * f[..., 57] * vy * x * y * z - + 1.590990257669731 * f[..., 56] * vx * x * y * z - + 0.9185586535436913 * f[..., 21] * x * y * z - + 1.590990257669731 * f[..., 60] * vx * vy * y * z - + 0.9185586535436913 * f[..., 27] * vy * y * z - + 0.9185586535436913 * f[..., 24] * vx * y * z - + 0.5303300858899105 * f[..., 8] * y * z - + 1.590990257669731 * f[..., 59] * vx * vy * x * z - + 0.9185586535436913 * f[..., 26] * vy * x * z - + 0.9185586535436913 * f[..., 23] * vx * x * z - + 0.5303300858899105 * f[..., 7] * x * z - + 0.9185586535436913 * f[..., 30] * vx * vy * z - + 0.5303300858899105 * f[..., 14] * vy * z - + 0.5303300858899105 * f[..., 11] * vx * z - + 0.3061862178478971 * f[..., 3] * z - + 6.075694404757366 * f[..., 175] * (vx * vy * x * y ^ 3 - 0.6 * vx * vy * x * y) - + 3.507803800100568 * f[..., 145] * (vy * x * y ^ 3 - 0.6 * vy * x * y) - + 3.507803800100568 * f[..., 136] * (vx * x * y ^ 3 - 0.6 * vx * x * y) - + 2.025231468252455 * f[..., 92] * (x * y ^ 3 - 0.6 * x * y) - + 3.507803800100568 * f[..., 151] * (vx * vy * y ^ 3 - 0.6 * vx * vy * y) - + 2.025231468252455 * f[..., 104] * (vy * y ^ 3 - 0.6 * vy * y) - + 2.025231468252455 * f[..., 98] * (vx * y ^ 3 - 0.6 * vx * y) - + 1.169267933366856 * f[..., 52] * (y ^ 3 - 0.6 * y) - + 3.080939385966558 - * f[..., 120] - * (vx * vy * x * y ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.778781183844713 * f[..., 74] * (vy * x * y ^ 2 - 0.3333333333333333 * vy * x) - + 1.778781183844713 * f[..., 65] * (vx * x * y ^ 2 - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 32] * (x * y ^ 2 - 0.3333333333333333 * x) - + 1.778781183844713 - * f[..., 80] - * (vx * vy * y ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 44] * (vy * y ^ 2 - 0.3333333333333333 * vy) - + 1.026979795322186 * f[..., 38] * (vx * y ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 17] * (y ^ 2 - 0.3333333333333333) - + 6.075694404757366 * f[..., 174] * (vx * vy * x ^ 3 * y - 0.6 * vx * vy * x * y) - + 3.507803800100568 * f[..., 144] * (vy * x ^ 3 * y - 0.6 * vy * x * y) - + 3.507803800100568 * f[..., 135] * (vx * x ^ 3 * y - 0.6 * vx * x * y) - + 2.025231468252455 * f[..., 91] * (x ^ 3 * y - 0.6 * x * y) - + 3.080939385966558 - * f[..., 119] - * (vx * vy * x ^ 2 * y - 0.3333333333333333 * vx * vy * y) - + 1.778781183844713 * f[..., 73] * (vy * x ^ 2 * y - 0.3333333333333333 * vy * y) - + 1.778781183844713 * f[..., 64] * (vx * x ^ 2 * y - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 31] * (x ^ 2 * y - 0.3333333333333333 * y) - + 6.075694404757366 * f[..., 184] * (vx * vy ^ 3 * x * y - 0.6 * vx * vy * x * y) - + 3.507803800100568 * f[..., 156] * (vy ^ 3 * x * y - 0.6 * vy * x * y) - + 3.080939385966558 - * f[..., 129] - * (vx * vy ^ 2 * x * y - 0.3333333333333333 * vx * x * y) - + 1.778781183844713 * f[..., 85] * (vy ^ 2 * x * y - 0.3333333333333333 * x * y) - + 6.075694404757366 * f[..., 180] * (vx ^ 3 * vy * x * y - 0.6 * vx * vy * x * y) - + 3.080939385966558 - * f[..., 125] - * (vx ^ 2 * vy * x * y - 0.3333333333333333 * vy * x * y) - + 3.507803800100568 * f[..., 141] * (vx ^ 3 * x * y - 0.6 * vx * x * y) - + 1.778781183844713 * f[..., 70] * (vx ^ 2 * x * y - 0.3333333333333333 * x * y) - + 3.507803800100568 * f[..., 160] * (vx * vy ^ 3 * y - 0.6 * vx * vy * y) - + 2.025231468252455 * f[..., 108] * (vy ^ 3 * y - 0.6 * vy * y) - + 1.778781183844713 * f[..., 89] * (vx * vy ^ 2 * y - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 48] * (vy ^ 2 * y - 0.3333333333333333 * y) - + 3.507803800100568 * f[..., 154] * (vx ^ 3 * vy * y - 0.6 * vx * vy * y) - + 1.778781183844713 * f[..., 83] * (vx ^ 2 * vy * y - 0.3333333333333333 * vy * y) - + 2.025231468252455 * f[..., 101] * (vx ^ 3 * y - 0.6 * vx * y) - + 1.026979795322186 * f[..., 41] * (vx ^ 2 * y - 0.3333333333333333 * y) - + 1.590990257669731 * f[..., 58] * vx * vy * x * y - + 0.9185586535436913 * f[..., 25] * vy * x * y - + 0.9185586535436913 * f[..., 22] * vx * x * y - + 0.5303300858899105 * f[..., 6] * x * y - + 0.9185586535436913 * f[..., 29] * vx * vy * y - + 0.5303300858899105 * f[..., 13] * vy * y - + 0.5303300858899105 * f[..., 10] * vx * y - + 0.3061862178478971 * f[..., 2] * y - + 3.507803800100568 * f[..., 150] * (vx * vy * x ^ 3 - 0.6 * vx * vy * x) - + 2.025231468252455 * f[..., 103] * (vy * x ^ 3 - 0.6 * vy * x) - + 2.025231468252455 * f[..., 97] * (vx * x ^ 3 - 0.6 * vx * x) - + 1.169267933366856 * f[..., 51] * (x ^ 3 - 0.6 * x) - + 1.778781183844713 - * f[..., 79] - * (vx * vy * x ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 43] * (vy * x ^ 2 - 0.3333333333333333 * vy) - + 1.026979795322186 * f[..., 37] * (vx * x ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 16] * (x ^ 2 - 0.3333333333333333) - + 3.507803800100568 * f[..., 159] * (vx * vy ^ 3 * x - 0.6 * vx * vy * x) - + 2.025231468252455 * f[..., 107] * (vy ^ 3 * x - 0.6 * vy * x) - + 1.778781183844713 * f[..., 88] * (vx * vy ^ 2 * x - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 47] * (vy ^ 2 * x - 0.3333333333333333 * x) - + 3.507803800100568 * f[..., 153] * (vx ^ 3 * vy * x - 0.6 * vx * vy * x) - + 1.778781183844713 * f[..., 82] * (vx ^ 2 * vy * x - 0.3333333333333333 * vy * x) - + 2.025231468252455 * f[..., 100] * (vx ^ 3 * x - 0.6 * vx * x) - + 1.026979795322186 * f[..., 40] * (vx ^ 2 * x - 0.3333333333333333 * x) - + 0.9185586535436913 * f[..., 28] * vx * vy * x - + 0.5303300858899105 * f[..., 12] * vy * x - + 0.5303300858899105 * f[..., 9] * vx * x - + 0.3061862178478971 * f[..., 1] * x - + 2.025231468252455 * f[..., 110] * (vx * vy ^ 3 - 0.6 * vx * vy) - + 1.169267933366856 * f[..., 55] * (vy ^ 3 - 0.6 * vy) - + 1.026979795322186 * f[..., 50] * (vx * vy ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 20] * (vy ^ 2 - 0.3333333333333333) - + 2.025231468252455 * f[..., 106] * (vx ^ 3 * vy - 0.6 * vx * vy) - + 1.026979795322186 * f[..., 46] * (vx ^ 2 * vy - 0.3333333333333333 * vy) - + 0.5303300858899105 * f[..., 15] * vx * vy - + 0.3061862178478971 * f[..., 5] * vy - + 1.169267933366856 * f[..., 54] * (vx ^ 3 - 0.6 * vx) - + 0.592927061281571 * f[..., 19] * (vx ^ 2 - 0.3333333333333333) - + 0.3061862178478971 * f[..., 4] * vx - + 0.1767766952966368 * f[..., 0] - ) - - -# end - - -def _expand_5d4p(f, x, y, z, vx, vy): - return ( - 20.88174713191522 - * f[..., 349] - * +12.05608232776094 - * f[..., 333] - * ( - vy * x * y * z - ^ 4 - - 0.8571428571428571 * (vy * x * y * z ^ 2 - 0.3333333333333333 * vy * x * y) - - 0.2 * vy * x * y - ) - + 12.05608232776094 - * f[..., 329] - * ( - vx * x * y * z - ^ 4 - - 0.8571428571428571 * (vx * x * y * z ^ 2 - 0.3333333333333333 * vx * x * y) - - 0.2 * vx * x * y - ) - + 6.960582377305072 - * f[..., 284] - * ( - x * y * z - ^ 4 - - 0.8571428571428571 * (x * y * z ^ 2 - 0.3333333333333333 * x * y) - - 0.2 * x * y - ) - + 12.05608232776094 - * f[..., 339] - * ( - vx * vy * y * z - ^ 4 - - 0.8571428571428571 - * (vx * vy * y * z ^ 2 - 0.3333333333333333 * vx * vy * y) - - 0.2 * vx * vy * y - ) - + 6.960582377305072 - * f[..., 299] - * ( - vy * y * z - ^ 4 - - 0.8571428571428571 * (vy * y * z ^ 2 - 0.3333333333333333 * vy * y) - - 0.2 * vy * y - ) - + 6.960582377305072 - * f[..., 290] - * ( - vx * y * z - ^ 4 - - 0.8571428571428571 * (vx * y * z ^ 2 - 0.3333333333333333 * vx * y) - - 0.2 * vx * y - ) - + 4.018694109253648 - * f[..., 212] - * ( - y * z - ^ 4 - 0.8571428571428571 * (y * z ^ 2 - 0.3333333333333333 * y) - 0.2 * y - ) - + 12.05608232776094 - * f[..., 338] - * ( - vx * vy * x * z - ^ 4 - - 0.8571428571428571 - * (vx * vy * x * z ^ 2 - 0.3333333333333333 * vx * vy * x) - - 0.2 * vx * vy * x - ) - + 6.960582377305072 - * f[..., 298] - * ( - vy * x * z - ^ 4 - - 0.8571428571428571 * (vy * x * z ^ 2 - 0.3333333333333333 * vy * x) - - 0.2 * vy * x - ) - + 6.960582377305072 - * f[..., 289] - * ( - vx * x * z - ^ 4 - - 0.8571428571428571 * (vx * x * z ^ 2 - 0.3333333333333333 * vx * x) - - 0.2 * vx * x - ) - + 4.018694109253648 - * f[..., 211] - * ( - x * z - ^ 4 - 0.8571428571428571 * (x * z ^ 2 - 0.3333333333333333 * x) - 0.2 * x - ) - + 6.960582377305072 - * f[..., 302] - * ( - vx * vy * z - ^ 4 - - 0.8571428571428571 * (vx * vy * z ^ 2 - 0.3333333333333333 * vx * vy) - - 0.2 * vx * vy - ) - + 4.018694109253648 - * f[..., 221] - * ( - vy * z - ^ 4 - 0.8571428571428571 * (vy * z ^ 2 - 0.3333333333333333 * vy) - 0.2 * vy - ) - + 4.018694109253648 - * f[..., 215] - * ( - vx * z - ^ 4 - 0.8571428571428571 * (vx * z ^ 2 - 0.3333333333333333 * vx) - 0.2 * vx - ) - + 2.320194125768357 - * f[..., 123] - * (z ^ 4 - 0.8571428571428571 * (z ^ 2 - 0.3333333333333333) - 0.2) - + 10.52341140030171 - * f[..., 324] - * (vx * vy * x * y * z ^ 3 - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 268] * (vy * x * y * z ^ 3 - 0.6 * vy * x * y * z) - + 6.075694404757366 * f[..., 264] * (vx * x * y * z ^ 3 - 0.6 * vx * x * y * z) - + 3.507803800100568 * f[..., 179] * (x * y * z ^ 3 - 0.6 * x * y * z) - + 6.075694404757366 * f[..., 274] * (vx * vy * y * z ^ 3 - 0.6 * vx * vy * y * z) - + 3.507803800100568 * f[..., 194] * (vy * y * z ^ 3 - 0.6 * vy * y * z) - + 3.507803800100568 * f[..., 185] * (vx * y * z ^ 3 - 0.6 * vx * y * z) - + 2.025231468252455 * f[..., 106] * (y * z ^ 3 - 0.6 * y * z) - + 6.075694404757366 * f[..., 273] * (vx * vy * x * z ^ 3 - 0.6 * vx * vy * x * z) - + 3.507803800100568 * f[..., 193] * (vy * x * z ^ 3 - 0.6 * vy * x * z) - + 3.507803800100568 * f[..., 184] * (vx * x * z ^ 3 - 0.6 * vx * x * z) - + 2.025231468252455 * f[..., 105] * (x * z ^ 3 - 0.6 * x * z) - + 3.507803800100568 * f[..., 197] * (vx * vy * z ^ 3 - 0.6 * vx * vy * z) - + 2.025231468252455 * f[..., 115] * (vy * z ^ 3 - 0.6 * vy * z) - + 2.025231468252455 * f[..., 109] * (vx * z ^ 3 - 0.6 * vx * z) - + 1.169267933366856 * f[..., 53] * (z ^ 3 - 0.6 * z) - + 10.33378485236652 - * f[..., 317] - * +10.33378485236652 - * f[..., 320] - * +5.966213466261491 - * f[..., 252] - * +5.966213466261491 - * f[..., 237] - * +10.33378485236652 - * f[..., 313] - * +5.966213466261491 - * f[..., 239] - * +5.966213466261491 - * f[..., 249] - * +5.966213466261491 - * f[..., 233] - * +5.966213466261491 - * f[..., 258] - * +3.444594950788841 - * f[..., 148] - * +3.444594950788841 - * f[..., 170] - * +3.444594950788841 - * f[..., 158] - * +10.33378485236652 - * f[..., 314] - * +5.966213466261491 - * f[..., 240] - * +5.966213466261491 - * f[..., 248] - * +5.966213466261491 - * f[..., 234] - * +5.966213466261491 - * f[..., 257] - * +3.444594950788841 - * f[..., 149] - * +3.444594950788841 - * f[..., 169] - * +3.444594950788841 - * f[..., 157] - * +5.966213466261491 - * f[..., 243] - * +5.966213466261491 - * f[..., 242] - * +3.444594950788841 - * f[..., 161] - * +3.444594950788841 - * f[..., 160] - * +3.444594950788841 - * f[..., 164] - * +3.444594950788841 - * f[..., 152] - * +3.444594950788841 - * f[..., 151] - * +3.444594950788841 - * f[..., 173] - * +1.988737822087164 - * f[..., 93] - * +5.336343551534138 - * f[..., 229] - * (vx * vy * x * y * z ^ 2 - 0.3333333333333333 * vx * vy * x * y) - + 3.080939385966558 - * f[..., 133] - * (vy * x * y * z ^ 2 - 0.3333333333333333 * vy * x * y) - + 3.080939385966558 - * f[..., 129] - * (vx * x * y * z ^ 2 - 0.3333333333333333 * vx * x * y) - + 1.778781183844713 * f[..., 63] * (x * y * z ^ 2 - 0.3333333333333333 * x * y) - + 3.080939385966558 - * f[..., 139] - * (vx * vy * y * z ^ 2 - 0.3333333333333333 * vx * vy * y) - + 1.778781183844713 * f[..., 78] * (vy * y * z ^ 2 - 0.3333333333333333 * vy * y) - + 1.778781183844713 * f[..., 69] * (vx * y * z ^ 2 - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 36] * (y * z ^ 2 - 0.3333333333333333 * y) - + 1.988737822087164 - * f[..., 92] - * +3.080939385966558 - * f[..., 138] - * (vx * vy * x * z ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.778781183844713 * f[..., 77] * (vy * x * z ^ 2 - 0.3333333333333333 * vy * x) - + 1.778781183844713 * f[..., 68] * (vx * x * z ^ 2 - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 35] * (x * z ^ 2 - 0.3333333333333333 * x) - + 1.988737822087164 - * f[..., 99] - * +1.778781183844713 - * f[..., 81] - * (vx * vy * z ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 45] * (vy * z ^ 2 - 0.3333333333333333 * vy) - + 1.988737822087164 - * f[..., 96] - * +1.026979795322186 - * f[..., 39] - * (vx * z ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 18] * (z ^ 2 - 0.3333333333333333) - + 20.88174713191522 - * f[..., 348] - * +10.33378485236652 - * f[..., 316] - * +12.05608232776094 - * f[..., 332] - * ( - -0.8571428571428571 * (vy * x * y ^ 2 * z - 0.3333333333333333 * vy * x * z) - + vy * x * y - ^ 4 * z - 0.2 * vy * x * z - ) - + 10.33378485236652 - * f[..., 319] - * +12.05608232776094 - * f[..., 328] - * ( - -0.8571428571428571 * (vx * x * y ^ 2 * z - 0.3333333333333333 * vx * x * z) - + vx * x * y - ^ 4 * z - 0.2 * vx * x * z - ) - + 5.966213466261491 - * f[..., 251] - * +5.966213466261491 - * f[..., 236] - * +6.960582377305072 - * f[..., 283] - * ( - -0.8571428571428571 * (x * y ^ 2 * z - 0.3333333333333333 * x * z) + x * y - ^ 4 * z - 0.2 * x * z - ) - + 10.33378485236652 - * f[..., 312] - * +12.05608232776094 - * f[..., 337] - * ( - -0.8571428571428571 * (vx * vy * y ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + vx * vy * y - ^ 4 * z - 0.2 * vx * vy * z - ) - + 5.966213466261491 - * f[..., 238] - * +5.966213466261491 - * f[..., 247] - * +6.960582377305072 - * f[..., 297] - * ( - -0.8571428571428571 * (vy * y ^ 2 * z - 0.3333333333333333 * vy * z) + vy * y - ^ 4 * z - 0.2 * vy * z - ) - + 5.966213466261491 - * f[..., 232] - * +5.966213466261491 - * f[..., 256] - * +6.960582377305072 - * f[..., 288] - * ( - -0.8571428571428571 * (vx * y ^ 2 * z - 0.3333333333333333 * vx * z) + vx * y - ^ 4 * z - 0.2 * vx * z - ) - + 3.444594950788841 - * f[..., 147] - * +3.444594950788841 - * f[..., 168] - * +3.444594950788841 - * f[..., 156] - * +4.018694109253648 - * f[..., 210] - * ( - -0.8571428571428571 * (y ^ 2 * z - 0.3333333333333333 * z) + y - ^ 4 * z - 0.2 * z - ) - + 20.88174713191522 - * f[..., 347] - * +10.33378485236652 - * f[..., 315] - * +12.05608232776094 - * f[..., 331] - * ( - -0.8571428571428571 * (vy * x ^ 2 * y * z - 0.3333333333333333 * vy * y * z) - + vy * x - ^ 4 * y * z - 0.2 * vy * y * z - ) - + 10.33378485236652 - * f[..., 318] - * +12.05608232776094 - * f[..., 327] - * ( - -0.8571428571428571 * (vx * x ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + vx * x - ^ 4 * y * z - 0.2 * vx * y * z - ) - + 5.966213466261491 - * f[..., 250] - * +5.966213466261491 - * f[..., 235] - * +6.960582377305072 - * f[..., 282] - * ( - -0.8571428571428571 * (x ^ 2 * y * z - 0.3333333333333333 * y * z) + x - ^ 4 * y * z - 0.2 * y * z - ) - + 20.88174713191522 - * f[..., 351] - * +10.33378485236652 - * f[..., 321] - * +12.05608232776094 - * f[..., 343] - * ( - -0.8571428571428571 * (vy ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + vy - ^ 4 * x * y * z - 0.2 * x * y * z - ) - + 20.88174713191522 - * f[..., 350] - * +12.05608232776094 - * f[..., 330] - * ( - -0.8571428571428571 * (vx ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + vx - ^ 4 * x * y * z - 0.2 * x * y * z - ) - + 12.05608232776094 - * f[..., 346] - * ( - -0.8571428571428571 * (vx * vy ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + vx * vy - ^ 4 * y * z - 0.2 * vx * y * z - ) - + 5.966213466261491 - * f[..., 261] - * +6.960582377305072 - * f[..., 308] - * ( - -0.8571428571428571 * (vy ^ 2 * y * z - 0.3333333333333333 * y * z) + vy - ^ 4 * y * z - 0.2 * y * z - ) - + 12.05608232776094 - * f[..., 342] - * ( - -0.8571428571428571 * (vx ^ 2 * vy * y * z - 0.3333333333333333 * vy * y * z) - + vx - ^ 4 * vy * y * z - 0.2 * vy * y * z - ) - + 6.960582377305072 - * f[..., 293] - * ( - -0.8571428571428571 * (vx ^ 2 * y * z - 0.3333333333333333 * y * z) + vx - ^ 4 * y * z - 0.2 * y * z - ) - + 12.05608232776094 - * f[..., 336] - * ( - -0.8571428571428571 * (vx * vy * x ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + vx * vy * x - ^ 4 * z - 0.2 * vx * vy * z - ) - + 5.966213466261491 - * f[..., 246] - * +6.960582377305072 - * f[..., 296] - * ( - -0.8571428571428571 * (vy * x ^ 2 * z - 0.3333333333333333 * vy * z) + vy * x - ^ 4 * z - 0.2 * vy * z - ) - + 5.966213466261491 - * f[..., 255] - * +6.960582377305072 - * f[..., 287] - * ( - -0.8571428571428571 * (vx * x ^ 2 * z - 0.3333333333333333 * vx * z) + vx * x - ^ 4 * z - 0.2 * vx * z - ) - + 3.444594950788841 - * f[..., 167] - * +3.444594950788841 - * f[..., 155] - * +4.018694109253648 - * f[..., 209] - * ( - -0.8571428571428571 * (x ^ 2 * z - 0.3333333333333333 * z) + x - ^ 4 * z - 0.2 * z - ) - + 12.05608232776094 - * f[..., 345] - * ( - -0.8571428571428571 * (vx * vy ^ 2 * x * z - 0.3333333333333333 * vx * x * z) - + vx * vy - ^ 4 * x * z - 0.2 * vx * x * z - ) - + 5.966213466261491 - * f[..., 260] - * +6.960582377305072 - * f[..., 307] - * ( - -0.8571428571428571 * (vy ^ 2 * x * z - 0.3333333333333333 * x * z) + vy - ^ 4 * x * z - 0.2 * x * z - ) - + 12.05608232776094 - * f[..., 341] - * ( - -0.8571428571428571 * (vx ^ 2 * vy * x * z - 0.3333333333333333 * vy * x * z) - + vx - ^ 4 * vy * x * z - 0.2 * vy * x * z - ) - + 6.960582377305072 - * f[..., 292] - * ( - -0.8571428571428571 * (vx ^ 2 * x * z - 0.3333333333333333 * x * z) + vx - ^ 4 * x * z - 0.2 * x * z - ) - + 6.960582377305072 - * f[..., 311] - * ( - -0.8571428571428571 * (vx * vy ^ 2 * z - 0.3333333333333333 * vx * z) - + vx * vy - ^ 4 * z - 0.2 * vx * z - ) - + 3.444594950788841 - * f[..., 176] - * +4.018694109253648 - * f[..., 225] - * ( - -0.8571428571428571 * (vy ^ 2 * z - 0.3333333333333333 * z) + vy - ^ 4 * z - 0.2 * z - ) - + 6.960582377305072 - * f[..., 305] - * ( - -0.8571428571428571 * (vx ^ 2 * vy * z - 0.3333333333333333 * vy * z) + vx - ^ 4 * vy * z - 0.2 * vy * z - ) - + 4.018694109253648 - * f[..., 218] - * ( - -0.8571428571428571 * (vx ^ 2 * z - 0.3333333333333333 * z) + vx - ^ 4 * z - 0.2 * z - ) - + 10.52341140030171 - * f[..., 323] - * (vx * vy * x * y ^ 3 * z - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 267] * (vy * x * y ^ 3 * z - 0.6 * vy * x * y * z) - + 6.075694404757366 * f[..., 263] * (vx * x * y ^ 3 * z - 0.6 * vx * x * y * z) - + 3.507803800100568 * f[..., 178] * (x * y ^ 3 * z - 0.6 * x * y * z) - + 6.075694404757366 * f[..., 272] * (vx * vy * y ^ 3 * z - 0.6 * vx * vy * y * z) - + 3.507803800100568 * f[..., 192] * (vy * y ^ 3 * z - 0.6 * vy * y * z) - + 3.507803800100568 * f[..., 183] * (vx * y ^ 3 * z - 0.6 * vx * y * z) - + 2.025231468252455 * f[..., 104] * (y ^ 3 * z - 0.6 * y * z) - + 5.336343551534138 - * f[..., 228] - * (vx * vy * x * y ^ 2 * z - 0.3333333333333333 * vx * vy * x * z) - + 3.080939385966558 - * f[..., 132] - * (vy * x * y ^ 2 * z - 0.3333333333333333 * vy * x * z) - + 3.080939385966558 - * f[..., 128] - * (vx * x * y ^ 2 * z - 0.3333333333333333 * vx * x * z) - + 1.778781183844713 * f[..., 62] * (x * y ^ 2 * z - 0.3333333333333333 * x * z) - + 3.080939385966558 - * f[..., 137] - * (vx * vy * y ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.778781183844713 * f[..., 76] * (vy * y ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.778781183844713 * f[..., 67] * (vx * y ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 34] * (y ^ 2 * z - 0.3333333333333333 * z) - + 10.52341140030171 - * f[..., 322] - * (vx * vy * x ^ 3 * y * z - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 266] * (vy * x ^ 3 * y * z - 0.6 * vy * x * y * z) - + 6.075694404757366 * f[..., 262] * (vx * x ^ 3 * y * z - 0.6 * vx * x * y * z) - + 3.507803800100568 * f[..., 177] * (x ^ 3 * y * z - 0.6 * x * y * z) - + 5.336343551534138 - * f[..., 227] - * (vx * vy * x ^ 2 * y * z - 0.3333333333333333 * vx * vy * y * z) - + 3.080939385966558 - * f[..., 131] - * (vy * x ^ 2 * y * z - 0.3333333333333333 * vy * y * z) - + 3.080939385966558 - * f[..., 127] - * (vx * x ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.778781183844713 * f[..., 61] * (x ^ 2 * y * z - 0.3333333333333333 * y * z) - + 10.52341140030171 - * f[..., 326] - * (vx * vy ^ 3 * x * y * z - 0.6 * vx * vy * x * y * z) - + 6.075694404757366 * f[..., 278] * (vy ^ 3 * x * y * z - 0.6 * vy * x * y * z) - + 5.336343551534138 - * f[..., 231] - * (vx * vy ^ 2 * x * y * z - 0.3333333333333333 * vx * x * y * z) - + 3.080939385966558 - * f[..., 143] - * (vy ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 10.52341140030171 - * f[..., 325] - * (vx ^ 3 * vy * x * y * z - 0.6 * vx * vy * x * y * z) - + 5.336343551534138 - * f[..., 230] - * (vx ^ 2 * vy * x * y * z - 0.3333333333333333 * vy * x * y * z) - + 6.075694404757366 * f[..., 265] * (vx ^ 3 * x * y * z - 0.6 * vx * x * y * z) - + 3.080939385966558 - * f[..., 130] - * (vx ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 6.075694404757366 * f[..., 281] * (vx * vy ^ 3 * y * z - 0.6 * vx * vy * y * z) - + 3.507803800100568 * f[..., 203] * (vy ^ 3 * y * z - 0.6 * vy * y * z) - + 3.080939385966558 - * f[..., 146] - * (vx * vy ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.778781183844713 * f[..., 87] * (vy ^ 2 * y * z - 0.3333333333333333 * y * z) - + 6.075694404757366 * f[..., 277] * (vx ^ 3 * vy * y * z - 0.6 * vx * vy * y * z) - + 3.080939385966558 - * f[..., 142] - * (vx ^ 2 * vy * y * z - 0.3333333333333333 * vy * y * z) - + 3.507803800100568 * f[..., 188] * (vx ^ 3 * y * z - 0.6 * vx * y * z) - + 1.778781183844713 * f[..., 72] * (vx ^ 2 * y * z - 0.3333333333333333 * y * z) - + 6.075694404757366 * f[..., 271] * (vx * vy * x ^ 3 * z - 0.6 * vx * vy * x * z) - + 3.507803800100568 * f[..., 191] * (vy * x ^ 3 * z - 0.6 * vy * x * z) - + 3.507803800100568 * f[..., 182] * (vx * x ^ 3 * z - 0.6 * vx * x * z) - + 2.025231468252455 * f[..., 103] * (x ^ 3 * z - 0.6 * x * z) - + 3.080939385966558 - * f[..., 136] - * (vx * vy * x ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.778781183844713 * f[..., 75] * (vy * x ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.778781183844713 * f[..., 66] * (vx * x ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 33] * (x ^ 2 * z - 0.3333333333333333 * z) - + 6.075694404757366 * f[..., 280] * (vx * vy ^ 3 * x * z - 0.6 * vx * vy * x * z) - + 3.507803800100568 * f[..., 202] * (vy ^ 3 * x * z - 0.6 * vy * x * z) - + 3.080939385966558 - * f[..., 145] - * (vx * vy ^ 2 * x * z - 0.3333333333333333 * vx * x * z) - + 1.778781183844713 * f[..., 86] * (vy ^ 2 * x * z - 0.3333333333333333 * x * z) - + 6.075694404757366 * f[..., 276] * (vx ^ 3 * vy * x * z - 0.6 * vx * vy * x * z) - + 3.080939385966558 - * f[..., 141] - * (vx ^ 2 * vy * x * z - 0.3333333333333333 * vy * x * z) - + 3.507803800100568 * f[..., 187] * (vx ^ 3 * x * z - 0.6 * vx * x * z) - + 1.778781183844713 * f[..., 71] * (vx ^ 2 * x * z - 0.3333333333333333 * x * z) - + 3.507803800100568 * f[..., 206] * (vx * vy ^ 3 * z - 0.6 * vx * vy * z) - + 2.025231468252455 * f[..., 119] * (vy ^ 3 * z - 0.6 * vy * z) - + 1.778781183844713 * f[..., 90] * (vx * vy ^ 2 * z - 0.3333333333333333 * vx * z) - + 1.026979795322186 * f[..., 49] * (vy ^ 2 * z - 0.3333333333333333 * z) - + 3.507803800100568 * f[..., 200] * (vx ^ 3 * vy * z - 0.6 * vx * vy * z) - + 1.778781183844713 * f[..., 84] * (vx ^ 2 * vy * z - 0.3333333333333333 * vy * z) - + 2.025231468252455 * f[..., 112] * (vx ^ 3 * z - 0.6 * vx * z) - + 1.026979795322186 * f[..., 42] * (vx ^ 2 * z - 0.3333333333333333 * z) - + 2.755675960631073 * f[..., 126] * vx * vy * x * y * z - + 1.590990257669731 * f[..., 57] * vy * x * y * z - + 1.590990257669731 * f[..., 56] * vx * x * y * z - + 0.9185586535436913 * f[..., 21] * x * y * z - + 1.590990257669731 * f[..., 60] * vx * vy * y * z - + 0.9185586535436913 * f[..., 27] * vy * y * z - + 0.9185586535436913 * f[..., 24] * vx * y * z - + 0.5303300858899105 * f[..., 8] * y * z - + 1.590990257669731 * f[..., 59] * vx * vy * x * z - + 0.9185586535436913 * f[..., 26] * vy * x * z - + 0.9185586535436913 * f[..., 23] * vx * x * z - + 0.5303300858899105 * f[..., 7] * x * z - + 0.9185586535436913 * f[..., 30] * vx * vy * z - + 0.5303300858899105 * f[..., 14] * vy * z - + 0.5303300858899105 * f[..., 11] * vx * z - + 0.3061862178478971 * f[..., 3] * z - + 12.05608232776094 - * f[..., 335] - * ( - vx * vy * x * y - ^ 4 - - 0.8571428571428571 - * (vx * vy * x * y ^ 2 - 0.3333333333333333 * vx * vy * x) - - 0.2 * vx * vy * x - ) - + 6.960582377305072 - * f[..., 295] - * ( - vy * x * y - ^ 4 - - 0.8571428571428571 * (vy * x * y ^ 2 - 0.3333333333333333 * vy * x) - - 0.2 * vy * x - ) - + 6.960582377305072 - * f[..., 286] - * ( - vx * x * y - ^ 4 - - 0.8571428571428571 * (vx * x * y ^ 2 - 0.3333333333333333 * vx * x) - - 0.2 * vx * x - ) - + 4.018694109253648 - * f[..., 208] - * ( - x * y - ^ 4 - 0.8571428571428571 * (x * y ^ 2 - 0.3333333333333333 * x) - 0.2 * x - ) - + 6.960582377305072 - * f[..., 301] - * ( - vx * vy * y - ^ 4 - - 0.8571428571428571 * (vx * vy * y ^ 2 - 0.3333333333333333 * vx * vy) - - 0.2 * vx * vy - ) - + 4.018694109253648 - * f[..., 220] - * ( - vy * y - ^ 4 - 0.8571428571428571 * (vy * y ^ 2 - 0.3333333333333333 * vy) - 0.2 * vy - ) - + 4.018694109253648 - * f[..., 214] - * ( - vx * y - ^ 4 - 0.8571428571428571 * (vx * y ^ 2 - 0.3333333333333333 * vx) - 0.2 * vx - ) - + 2.320194125768357 - * f[..., 122] - * (y ^ 4 - 0.8571428571428571 * (y ^ 2 - 0.3333333333333333) - 0.2) - + 6.075694404757366 * f[..., 270] * (vx * vy * x * y ^ 3 - 0.6 * vx * vy * x * y) - + 3.507803800100568 * f[..., 190] * (vy * x * y ^ 3 - 0.6 * vy * x * y) - + 3.507803800100568 * f[..., 181] * (vx * x * y ^ 3 - 0.6 * vx * x * y) - + 2.025231468252455 * f[..., 102] * (x * y ^ 3 - 0.6 * x * y) - + 3.507803800100568 * f[..., 196] * (vx * vy * y ^ 3 - 0.6 * vx * vy * y) - + 2.025231468252455 * f[..., 114] * (vy * y ^ 3 - 0.6 * vy * y) - + 2.025231468252455 * f[..., 108] * (vx * y ^ 3 - 0.6 * vx * y) - + 1.169267933366856 * f[..., 52] * (y ^ 3 - 0.6 * y) - + 5.966213466261491 - * f[..., 245] - * +5.966213466261491 - * f[..., 254] - * +3.444594950788841 - * f[..., 166] - * +3.444594950788841 - * f[..., 154] - * +5.966213466261491 - * f[..., 241] - * +3.444594950788841 - * f[..., 159] - * +3.444594950788841 - * f[..., 163] - * +3.444594950788841 - * f[..., 150] - * +3.444594950788841 - * f[..., 172] - * +1.988737822087164 - * f[..., 91] - * +3.080939385966558 - * f[..., 135] - * (vx * vy * x * y ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.778781183844713 * f[..., 74] * (vy * x * y ^ 2 - 0.3333333333333333 * vy * x) - + 1.778781183844713 * f[..., 65] * (vx * x * y ^ 2 - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 32] * (x * y ^ 2 - 0.3333333333333333 * x) - + 1.988737822087164 - * f[..., 98] - * +1.778781183844713 - * f[..., 80] - * (vx * vy * y ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 44] * (vy * y ^ 2 - 0.3333333333333333 * vy) - + 1.988737822087164 - * f[..., 95] - * +1.026979795322186 - * f[..., 38] - * (vx * y ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 17] * (y ^ 2 - 0.3333333333333333) - + 12.05608232776094 - * f[..., 334] - * ( - -0.8571428571428571 * (vx * vy * x ^ 2 * y - 0.3333333333333333 * vx * vy * y) - + vx * vy * x - ^ 4 * y - 0.2 * vx * vy * y - ) - + 5.966213466261491 - * f[..., 244] - * +6.960582377305072 - * f[..., 294] - * ( - -0.8571428571428571 * (vy * x ^ 2 * y - 0.3333333333333333 * vy * y) + vy * x - ^ 4 * y - 0.2 * vy * y - ) - + 5.966213466261491 - * f[..., 253] - * +6.960582377305072 - * f[..., 285] - * ( - -0.8571428571428571 * (vx * x ^ 2 * y - 0.3333333333333333 * vx * y) + vx * x - ^ 4 * y - 0.2 * vx * y - ) - + 3.444594950788841 - * f[..., 165] - * +3.444594950788841 - * f[..., 153] - * +4.018694109253648 - * f[..., 207] - * ( - -0.8571428571428571 * (x ^ 2 * y - 0.3333333333333333 * y) + x - ^ 4 * y - 0.2 * y - ) - + 12.05608232776094 - * f[..., 344] - * ( - -0.8571428571428571 * (vx * vy ^ 2 * x * y - 0.3333333333333333 * vx * x * y) - + vx * vy - ^ 4 * x * y - 0.2 * vx * x * y - ) - + 5.966213466261491 - * f[..., 259] - * +6.960582377305072 - * f[..., 306] - * ( - -0.8571428571428571 * (vy ^ 2 * x * y - 0.3333333333333333 * x * y) + vy - ^ 4 * x * y - 0.2 * x * y - ) - + 12.05608232776094 - * f[..., 340] - * ( - -0.8571428571428571 * (vx ^ 2 * vy * x * y - 0.3333333333333333 * vy * x * y) - + vx - ^ 4 * vy * x * y - 0.2 * vy * x * y - ) - + 6.960582377305072 - * f[..., 291] - * ( - -0.8571428571428571 * (vx ^ 2 * x * y - 0.3333333333333333 * x * y) + vx - ^ 4 * x * y - 0.2 * x * y - ) - + 6.960582377305072 - * f[..., 310] - * ( - -0.8571428571428571 * (vx * vy ^ 2 * y - 0.3333333333333333 * vx * y) - + vx * vy - ^ 4 * y - 0.2 * vx * y - ) - + 3.444594950788841 - * f[..., 175] - * +4.018694109253648 - * f[..., 224] - * ( - -0.8571428571428571 * (vy ^ 2 * y - 0.3333333333333333 * y) + vy - ^ 4 * y - 0.2 * y - ) - + 6.960582377305072 - * f[..., 304] - * ( - -0.8571428571428571 * (vx ^ 2 * vy * y - 0.3333333333333333 * vy * y) + vx - ^ 4 * vy * y - 0.2 * vy * y - ) - + 4.018694109253648 - * f[..., 217] - * ( - -0.8571428571428571 * (vx ^ 2 * y - 0.3333333333333333 * y) + vx - ^ 4 * y - 0.2 * y - ) - + 6.075694404757366 * f[..., 269] * (vx * vy * x ^ 3 * y - 0.6 * vx * vy * x * y) - + 3.507803800100568 * f[..., 189] * (vy * x ^ 3 * y - 0.6 * vy * x * y) - + 3.507803800100568 * f[..., 180] * (vx * x ^ 3 * y - 0.6 * vx * x * y) - + 2.025231468252455 * f[..., 101] * (x ^ 3 * y - 0.6 * x * y) - + 3.080939385966558 - * f[..., 134] - * (vx * vy * x ^ 2 * y - 0.3333333333333333 * vx * vy * y) - + 1.778781183844713 * f[..., 73] * (vy * x ^ 2 * y - 0.3333333333333333 * vy * y) - + 1.778781183844713 * f[..., 64] * (vx * x ^ 2 * y - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 31] * (x ^ 2 * y - 0.3333333333333333 * y) - + 6.075694404757366 * f[..., 279] * (vx * vy ^ 3 * x * y - 0.6 * vx * vy * x * y) - + 3.507803800100568 * f[..., 201] * (vy ^ 3 * x * y - 0.6 * vy * x * y) - + 3.080939385966558 - * f[..., 144] - * (vx * vy ^ 2 * x * y - 0.3333333333333333 * vx * x * y) - + 1.778781183844713 * f[..., 85] * (vy ^ 2 * x * y - 0.3333333333333333 * x * y) - + 6.075694404757366 * f[..., 275] * (vx ^ 3 * vy * x * y - 0.6 * vx * vy * x * y) - + 3.080939385966558 - * f[..., 140] - * (vx ^ 2 * vy * x * y - 0.3333333333333333 * vy * x * y) - + 3.507803800100568 * f[..., 186] * (vx ^ 3 * x * y - 0.6 * vx * x * y) - + 1.778781183844713 * f[..., 70] * (vx ^ 2 * x * y - 0.3333333333333333 * x * y) - + 3.507803800100568 * f[..., 205] * (vx * vy ^ 3 * y - 0.6 * vx * vy * y) - + 2.025231468252455 * f[..., 118] * (vy ^ 3 * y - 0.6 * vy * y) - + 1.778781183844713 * f[..., 89] * (vx * vy ^ 2 * y - 0.3333333333333333 * vx * y) - + 1.026979795322186 * f[..., 48] * (vy ^ 2 * y - 0.3333333333333333 * y) - + 3.507803800100568 * f[..., 199] * (vx ^ 3 * vy * y - 0.6 * vx * vy * y) - + 1.778781183844713 * f[..., 83] * (vx ^ 2 * vy * y - 0.3333333333333333 * vy * y) - + 2.025231468252455 * f[..., 111] * (vx ^ 3 * y - 0.6 * vx * y) - + 1.026979795322186 * f[..., 41] * (vx ^ 2 * y - 0.3333333333333333 * y) - + 1.590990257669731 * f[..., 58] * vx * vy * x * y - + 0.9185586535436913 * f[..., 25] * vy * x * y - + 0.9185586535436913 * f[..., 22] * vx * x * y - + 0.5303300858899105 * f[..., 6] * x * y - + 0.9185586535436913 * f[..., 29] * vx * vy * y - + 0.5303300858899105 * f[..., 13] * vy * y - + 0.5303300858899105 * f[..., 10] * vx * y - + 0.3061862178478971 * f[..., 2] * y - + 6.960582377305072 - * f[..., 300] - * ( - vx * vy * x - ^ 4 - - 0.8571428571428571 * (vx * vy * x ^ 2 - 0.3333333333333333 * vx * vy) - - 0.2 * vx * vy - ) - + 4.018694109253648 - * f[..., 219] - * ( - vy * x - ^ 4 - 0.8571428571428571 * (vy * x ^ 2 - 0.3333333333333333 * vy) - 0.2 * vy - ) - + 4.018694109253648 - * f[..., 213] - * ( - vx * x - ^ 4 - 0.8571428571428571 * (vx * x ^ 2 - 0.3333333333333333 * vx) - 0.2 * vx - ) - + 2.320194125768357 - * f[..., 121] - * (x ^ 4 - 0.8571428571428571 * (x ^ 2 - 0.3333333333333333) - 0.2) - + 3.507803800100568 * f[..., 195] * (vx * vy * x ^ 3 - 0.6 * vx * vy * x) - + 2.025231468252455 * f[..., 113] * (vy * x ^ 3 - 0.6 * vy * x) - + 2.025231468252455 * f[..., 107] * (vx * x ^ 3 - 0.6 * vx * x) - + 1.169267933366856 * f[..., 51] * (x ^ 3 - 0.6 * x) - + 3.444594950788841 - * f[..., 162] - * +3.444594950788841 - * f[..., 171] - * +1.988737822087164 - * f[..., 97] - * +1.778781183844713 - * f[..., 79] - * (vx * vy * x ^ 2 - 0.3333333333333333 * vx * vy) - + 1.026979795322186 * f[..., 43] * (vy * x ^ 2 - 0.3333333333333333 * vy) - + 1.988737822087164 - * f[..., 94] - * +1.026979795322186 - * f[..., 37] - * (vx * x ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 16] * (x ^ 2 - 0.3333333333333333) - + 6.960582377305072 - * f[..., 309] - * ( - -0.8571428571428571 * (vx * vy ^ 2 * x - 0.3333333333333333 * vx * x) - + vx * vy - ^ 4 * x - 0.2 * vx * x - ) - + 3.444594950788841 - * f[..., 174] - * +4.018694109253648 - * f[..., 223] - * ( - -0.8571428571428571 * (vy ^ 2 * x - 0.3333333333333333 * x) + vy - ^ 4 * x - 0.2 * x - ) - + 6.960582377305072 - * f[..., 303] - * ( - -0.8571428571428571 * (vx ^ 2 * vy * x - 0.3333333333333333 * vy * x) + vx - ^ 4 * vy * x - 0.2 * vy * x - ) - + 4.018694109253648 - * f[..., 216] - * ( - -0.8571428571428571 * (vx ^ 2 * x - 0.3333333333333333 * x) + vx - ^ 4 * x - 0.2 * x - ) - + 3.507803800100568 * f[..., 204] * (vx * vy ^ 3 * x - 0.6 * vx * vy * x) - + 2.025231468252455 * f[..., 117] * (vy ^ 3 * x - 0.6 * vy * x) - + 1.778781183844713 * f[..., 88] * (vx * vy ^ 2 * x - 0.3333333333333333 * vx * x) - + 1.026979795322186 * f[..., 47] * (vy ^ 2 * x - 0.3333333333333333 * x) - + 3.507803800100568 * f[..., 198] * (vx ^ 3 * vy * x - 0.6 * vx * vy * x) - + 1.778781183844713 * f[..., 82] * (vx ^ 2 * vy * x - 0.3333333333333333 * vy * x) - + 2.025231468252455 * f[..., 110] * (vx ^ 3 * x - 0.6 * vx * x) - + 1.026979795322186 * f[..., 40] * (vx ^ 2 * x - 0.3333333333333333 * x) - + 0.9185586535436913 * f[..., 28] * vx * vy * x - + 0.5303300858899105 * f[..., 12] * vy * x - + 0.5303300858899105 * f[..., 9] * vx * x - + 0.3061862178478971 * f[..., 1] * x - + 4.018694109253648 - * f[..., 226] - * ( - vx * vy - ^ 4 - 0.8571428571428571 * (vx * vy ^ 2 - 0.3333333333333333 * vx) - 0.2 * vx - ) - + 2.320194125768357 - * f[..., 125] - * (vy ^ 4 - 0.8571428571428571 * (vy ^ 2 - 0.3333333333333333) - 0.2) - + 2.025231468252455 * f[..., 120] * (vx * vy ^ 3 - 0.6 * vx * vy) - + 1.169267933366856 * f[..., 55] * (vy ^ 3 - 0.6 * vy) - + 1.988737822087164 - * f[..., 100] - * +1.026979795322186 - * f[..., 50] - * (vx * vy ^ 2 - 0.3333333333333333 * vx) - + 0.592927061281571 * f[..., 20] * (vy ^ 2 - 0.3333333333333333) - + 4.018694109253648 - * f[..., 222] - * ( - -0.8571428571428571 * (vx ^ 2 * vy - 0.3333333333333333 * vy) + vx - ^ 4 * vy - 0.2 * vy - ) - + 2.025231468252455 * f[..., 116] * (vx ^ 3 * vy - 0.6 * vx * vy) - + 1.026979795322186 * f[..., 46] * (vx ^ 2 * vy - 0.3333333333333333 * vy) - + 0.5303300858899105 * f[..., 15] * vx * vy - + 0.3061862178478971 * f[..., 5] * vy - + 2.320194125768357 - * f[..., 124] - * (vx ^ 4 - 0.8571428571428571 * (vx ^ 2 - 0.3333333333333333) - 0.2) - + 1.169267933366856 * f[..., 54] * (vx ^ 3 - 0.6 * vx) - + 0.592927061281571 * f[..., 19] * (vx ^ 2 - 0.3333333333333333) - + 0.3061862178478971 * f[..., 4] * vx - + 0.1767766952966368 * f[..., 0] - ) - - -# end - -expand_5d = [_expand_5d1p, _expand_5d2p, _expand_5d3p, _expand_5d4p] diff --git a/src/postgkyl/modalDG/kernels/expand6d.py b/src/postgkyl/modalDG/kernels/expand6d.py deleted file mode 100755 index 2e6b288b..00000000 --- a/src/postgkyl/modalDG/kernels/expand6d.py +++ /dev/null @@ -1,602 +0,0 @@ -def _expand_6d1p(f, x, y, z, vx, vy, vz): - return ( - 3.375 * f[..., 63] * vx * vy * vz * x * y * z - + 1.948557158514986 * f[..., 59] * vy * vz * x * y * z - + 1.948557158514986 * f[..., 58] * vx * vz * x * y * z - + 1.125 * f[..., 47] * vz * x * y * z - + 1.948557158514986 * f[..., 57] * vx * vy * x * y * z - + 1.125 * f[..., 43] * vy * x * y * z - + 1.125 * f[..., 42] * vx * x * y * z - + 0.6495190528383289 * f[..., 22] * x * y * z - + 1.948557158514986 * f[..., 62] * vx * vy * vz * y * z - + 1.125 * f[..., 53] * vy * vz * y * z - + 1.125 * f[..., 50] * vx * vz * y * z - + 0.6495190528383289 * f[..., 34] * vz * y * z - + 1.125 * f[..., 46] * vx * vy * y * z - + 0.6495190528383289 * f[..., 28] * vy * y * z - + 0.6495190528383289 * f[..., 25] * vx * y * z - + 0.375 * f[..., 9] * y * z - + 1.948557158514986 * f[..., 61] * vx * vy * vz * x * z - + 1.125 * f[..., 52] * vy * vz * x * z - + 1.125 * f[..., 49] * vx * vz * x * z - + 0.6495190528383289 * f[..., 33] * vz * x * z - + 1.125 * f[..., 45] * vx * vy * x * z - + 0.6495190528383289 * f[..., 27] * vy * x * z - + 0.6495190528383289 * f[..., 24] * vx * x * z - + 0.375 * f[..., 8] * x * z - + 1.125 * f[..., 56] * vx * vy * vz * z - + 0.6495190528383289 * f[..., 40] * vy * vz * z - + 0.6495190528383289 * f[..., 37] * vx * vz * z - + 0.375 * f[..., 19] * vz * z - + 0.6495190528383289 * f[..., 31] * vx * vy * z - + 0.375 * f[..., 15] * vy * z - + 0.375 * f[..., 12] * vx * z - + 0.2165063509461096 * f[..., 3] * z - + 1.948557158514986 * f[..., 60] * vx * vy * vz * x * y - + 1.125 * f[..., 51] * vy * vz * x * y - + 1.125 * f[..., 48] * vx * vz * x * y - + 0.6495190528383289 * f[..., 32] * vz * x * y - + 1.125 * f[..., 44] * vx * vy * x * y - + 0.6495190528383289 * f[..., 26] * vy * x * y - + 0.6495190528383289 * f[..., 23] * vx * x * y - + 0.375 * f[..., 7] * x * y - + 1.125 * f[..., 55] * vx * vy * vz * y - + 0.6495190528383289 * f[..., 39] * vy * vz * y - + 0.6495190528383289 * f[..., 36] * vx * vz * y - + 0.375 * f[..., 18] * vz * y - + 0.6495190528383289 * f[..., 30] * vx * vy * y - + 0.375 * f[..., 14] * vy * y - + 0.375 * f[..., 11] * vx * y - + 0.2165063509461096 * f[..., 2] * y - + 1.125 * f[..., 54] * vx * vy * vz * x - + 0.6495190528383289 * f[..., 38] * vy * vz * x - + 0.6495190528383289 * f[..., 35] * vx * vz * x - + 0.375 * f[..., 17] * vz * x - + 0.6495190528383289 * f[..., 29] * vx * vy * x - + 0.375 * f[..., 13] * vy * x - + 0.375 * f[..., 10] * vx * x - + 0.2165063509461096 * f[..., 1] * x - + 0.6495190528383289 * f[..., 41] * vx * vy * vz - + 0.375 * f[..., 21] * vy * vz - + 0.375 * f[..., 20] * vx * vz - + 0.2165063509461096 * f[..., 6] * vz - + 0.375 * f[..., 16] * vx * vy - + 0.2165063509461096 * f[..., 5] * vy - + 0.2165063509461096 * f[..., 4] * vx - + 0.125 * f[..., 0] - ) - - -# end - - -def _expand_6d2p(f, x, y, z, vx, vy, vz): - return ( - 6.535659396725016 - * f[..., 252] - * (vx * vy * vz * x * y * z ^ 2 - 0.3333333333333333 * vx * vy * vz * x * y) - + 3.773364712030896 - * f[..., 231] - * (vy * vz * x * y * z ^ 2 - 0.3333333333333333 * vy * vz * x * y) - + 3.773364712030896 - * f[..., 227] - * (vx * vz * x * y * z ^ 2 - 0.3333333333333333 * vx * vz * x * y) - + 2.178553132241672 - * f[..., 181] - * (vz * x * y * z ^ 2 - 0.3333333333333333 * vz * x * y) - + 3.773364712030896 - * f[..., 222] - * (vx * vy * x * y * z ^ 2 - 0.3333333333333333 * vx * vy * x * y) - + 2.178553132241672 - * f[..., 165] - * (vy * x * y * z ^ 2 - 0.3333333333333333 * vy * x * y) - + 2.178553132241672 - * f[..., 161] - * (vx * x * y * z ^ 2 - 0.3333333333333333 * vx * x * y) - + 1.257788237343632 * f[..., 95] * (x * y * z ^ 2 - 0.3333333333333333 * x * y) - + 3.773364712030896 - * f[..., 237] - * (vx * vy * vz * y * z ^ 2 - 0.3333333333333333 * vx * vy * vz * y) - + 2.178553132241672 - * f[..., 196] - * (vy * vz * y * z ^ 2 - 0.3333333333333333 * vy * vz * y) - + 2.178553132241672 - * f[..., 187] - * (vx * vz * y * z ^ 2 - 0.3333333333333333 * vx * vz * y) - + 1.257788237343632 * f[..., 128] * (vz * y * z ^ 2 - 0.3333333333333333 * vz * y) - + 2.178553132241672 - * f[..., 171] - * (vx * vy * y * z ^ 2 - 0.3333333333333333 * vx * vy * y) - + 1.257788237343632 * f[..., 110] * (vy * y * z ^ 2 - 0.3333333333333333 * vy * y) - + 1.257788237343632 * f[..., 101] * (vx * y * z ^ 2 - 0.3333333333333333 * vx * y) - + 0.7261843774138907 * f[..., 53] * (y * z ^ 2 - 0.3333333333333333 * y) - + 3.773364712030896 - * f[..., 236] - * (vx * vy * vz * x * z ^ 2 - 0.3333333333333333 * vx * vy * vz * x) - + 2.178553132241672 - * f[..., 195] - * (vy * vz * x * z ^ 2 - 0.3333333333333333 * vy * vz * x) - + 2.178553132241672 - * f[..., 186] - * (vx * vz * x * z ^ 2 - 0.3333333333333333 * vx * vz * x) - + 1.257788237343632 * f[..., 127] * (vz * x * z ^ 2 - 0.3333333333333333 * vz * x) - + 2.178553132241672 - * f[..., 170] - * (vx * vy * x * z ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.257788237343632 * f[..., 109] * (vy * x * z ^ 2 - 0.3333333333333333 * vy * x) - + 1.257788237343632 * f[..., 100] * (vx * x * z ^ 2 - 0.3333333333333333 * vx * x) - + 0.7261843774138907 * f[..., 52] * (x * z ^ 2 - 0.3333333333333333 * x) - + 2.178553132241672 - * f[..., 199] - * (vx * vy * vz * z ^ 2 - 0.3333333333333333 * vx * vy * vz) - + 1.257788237343632 - * f[..., 137] - * (vy * vz * z ^ 2 - 0.3333333333333333 * vy * vz) - + 1.257788237343632 - * f[..., 131] - * (vx * vz * z ^ 2 - 0.3333333333333333 * vx * vz) - + 0.7261843774138907 * f[..., 70] * (vz * z ^ 2 - 0.3333333333333333 * vz) - + 1.257788237343632 - * f[..., 113] - * (vx * vy * z ^ 2 - 0.3333333333333333 * vx * vy) - + 0.7261843774138907 * f[..., 62] * (vy * z ^ 2 - 0.3333333333333333 * vy) - + 0.7261843774138907 * f[..., 56] * (vx * z ^ 2 - 0.3333333333333333 * vx) - + 0.4192627457812106 * f[..., 24] * (z ^ 2 - 0.3333333333333333) - + 6.535659396725016 - * f[..., 251] - * (vx * vy * vz * x * y ^ 2 * z - 0.3333333333333333 * vx * vy * vz * x * z) - + 3.773364712030896 - * f[..., 230] - * (vy * vz * x * y ^ 2 * z - 0.3333333333333333 * vy * vz * x * z) - + 3.773364712030896 - * f[..., 226] - * (vx * vz * x * y ^ 2 * z - 0.3333333333333333 * vx * vz * x * z) - + 2.178553132241672 - * f[..., 180] - * (vz * x * y ^ 2 * z - 0.3333333333333333 * vz * x * z) - + 3.773364712030896 - * f[..., 221] - * (vx * vy * x * y ^ 2 * z - 0.3333333333333333 * vx * vy * x * z) - + 2.178553132241672 - * f[..., 164] - * (vy * x * y ^ 2 * z - 0.3333333333333333 * vy * x * z) - + 2.178553132241672 - * f[..., 160] - * (vx * x * y ^ 2 * z - 0.3333333333333333 * vx * x * z) - + 1.257788237343632 * f[..., 94] * (x * y ^ 2 * z - 0.3333333333333333 * x * z) - + 3.773364712030896 - * f[..., 235] - * (vx * vy * vz * y ^ 2 * z - 0.3333333333333333 * vx * vy * vz * z) - + 2.178553132241672 - * f[..., 194] - * (vy * vz * y ^ 2 * z - 0.3333333333333333 * vy * vz * z) - + 2.178553132241672 - * f[..., 185] - * (vx * vz * y ^ 2 * z - 0.3333333333333333 * vx * vz * z) - + 1.257788237343632 * f[..., 126] * (vz * y ^ 2 * z - 0.3333333333333333 * vz * z) - + 2.178553132241672 - * f[..., 169] - * (vx * vy * y ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.257788237343632 * f[..., 108] * (vy * y ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.257788237343632 * f[..., 99] * (vx * y ^ 2 * z - 0.3333333333333333 * vx * z) - + 0.7261843774138907 * f[..., 51] * (y ^ 2 * z - 0.3333333333333333 * z) - + 6.535659396725016 - * f[..., 250] - * (vx * vy * vz * x ^ 2 * y * z - 0.3333333333333333 * vx * vy * vz * y * z) - + 3.773364712030896 - * f[..., 229] - * (vy * vz * x ^ 2 * y * z - 0.3333333333333333 * vy * vz * y * z) - + 3.773364712030896 - * f[..., 225] - * (vx * vz * x ^ 2 * y * z - 0.3333333333333333 * vx * vz * y * z) - + 2.178553132241672 - * f[..., 179] - * (vz * x ^ 2 * y * z - 0.3333333333333333 * vz * y * z) - + 3.773364712030896 - * f[..., 220] - * (vx * vy * x ^ 2 * y * z - 0.3333333333333333 * vx * vy * y * z) - + 2.178553132241672 - * f[..., 163] - * (vy * x ^ 2 * y * z - 0.3333333333333333 * vy * y * z) - + 2.178553132241672 - * f[..., 159] - * (vx * x ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.257788237343632 * f[..., 93] * (x ^ 2 * y * z - 0.3333333333333333 * y * z) - + 6.535659396725016 - * f[..., 255] - * (vx * vy * vz ^ 2 * x * y * z - 0.3333333333333333 * vx * vy * x * y * z) - + 3.773364712030896 - * f[..., 246] - * (vy * vz ^ 2 * x * y * z - 0.3333333333333333 * vy * x * y * z) - + 3.773364712030896 - * f[..., 245] - * (vx * vz ^ 2 * x * y * z - 0.3333333333333333 * vx * x * y * z) - + 2.178553132241672 - * f[..., 209] - * (vz ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 6.535659396725016 - * f[..., 254] - * (vx * vy ^ 2 * vz * x * y * z - 0.3333333333333333 * vx * vz * x * y * z) - + 3.773364712030896 - * f[..., 241] - * (vy ^ 2 * vz * x * y * z - 0.3333333333333333 * vz * x * y * z) - + 6.535659396725016 - * f[..., 253] - * (vx ^ 2 * vy * vz * x * y * z - 0.3333333333333333 * vy * vz * x * y * z) - + 3.773364712030896 - * f[..., 228] - * (vx ^ 2 * vz * x * y * z - 0.3333333333333333 * vz * x * y * z) - + 3.773364712030896 - * f[..., 224] - * (vx * vy ^ 2 * x * y * z - 0.3333333333333333 * vx * x * y * z) - + 2.178553132241672 - * f[..., 175] - * (vy ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 3.773364712030896 - * f[..., 223] - * (vx ^ 2 * vy * x * y * z - 0.3333333333333333 * vy * x * y * z) - + 2.178553132241672 - * f[..., 162] - * (vx ^ 2 * x * y * z - 0.3333333333333333 * x * y * z) - + 3.773364712030896 - * f[..., 249] - * (vx * vy * vz ^ 2 * y * z - 0.3333333333333333 * vx * vy * y * z) - + 2.178553132241672 - * f[..., 215] - * (vy * vz ^ 2 * y * z - 0.3333333333333333 * vy * y * z) - + 2.178553132241672 - * f[..., 212] - * (vx * vz ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.257788237343632 * f[..., 145] * (vz ^ 2 * y * z - 0.3333333333333333 * y * z) - + 3.773364712030896 - * f[..., 244] - * (vx * vy ^ 2 * vz * y * z - 0.3333333333333333 * vx * vz * y * z) - + 2.178553132241672 - * f[..., 205] - * (vy ^ 2 * vz * y * z - 0.3333333333333333 * vz * y * z) - + 3.773364712030896 - * f[..., 240] - * (vx ^ 2 * vy * vz * y * z - 0.3333333333333333 * vy * vz * y * z) - + 2.178553132241672 - * f[..., 190] - * (vx ^ 2 * vz * y * z - 0.3333333333333333 * vz * y * z) - + 2.178553132241672 - * f[..., 178] - * (vx * vy ^ 2 * y * z - 0.3333333333333333 * vx * y * z) - + 1.257788237343632 * f[..., 119] * (vy ^ 2 * y * z - 0.3333333333333333 * y * z) - + 2.178553132241672 - * f[..., 174] - * (vx ^ 2 * vy * y * z - 0.3333333333333333 * vy * y * z) - + 1.257788237343632 * f[..., 104] * (vx ^ 2 * y * z - 0.3333333333333333 * y * z) - + 3.773364712030896 - * f[..., 234] - * (vx * vy * vz * x ^ 2 * z - 0.3333333333333333 * vx * vy * vz * z) - + 2.178553132241672 - * f[..., 193] - * (vy * vz * x ^ 2 * z - 0.3333333333333333 * vy * vz * z) - + 2.178553132241672 - * f[..., 184] - * (vx * vz * x ^ 2 * z - 0.3333333333333333 * vx * vz * z) - + 1.257788237343632 * f[..., 125] * (vz * x ^ 2 * z - 0.3333333333333333 * vz * z) - + 2.178553132241672 - * f[..., 168] - * (vx * vy * x ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.257788237343632 * f[..., 107] * (vy * x ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.257788237343632 * f[..., 98] * (vx * x ^ 2 * z - 0.3333333333333333 * vx * z) - + 0.7261843774138907 * f[..., 50] * (x ^ 2 * z - 0.3333333333333333 * z) - + 3.773364712030896 - * f[..., 248] - * (vx * vy * vz ^ 2 * x * z - 0.3333333333333333 * vx * vy * x * z) - + 2.178553132241672 - * f[..., 214] - * (vy * vz ^ 2 * x * z - 0.3333333333333333 * vy * x * z) - + 2.178553132241672 - * f[..., 211] - * (vx * vz ^ 2 * x * z - 0.3333333333333333 * vx * x * z) - + 1.257788237343632 * f[..., 144] * (vz ^ 2 * x * z - 0.3333333333333333 * x * z) - + 3.773364712030896 - * f[..., 243] - * (vx * vy ^ 2 * vz * x * z - 0.3333333333333333 * vx * vz * x * z) - + 2.178553132241672 - * f[..., 204] - * (vy ^ 2 * vz * x * z - 0.3333333333333333 * vz * x * z) - + 3.773364712030896 - * f[..., 239] - * (vx ^ 2 * vy * vz * x * z - 0.3333333333333333 * vy * vz * x * z) - + 2.178553132241672 - * f[..., 189] - * (vx ^ 2 * vz * x * z - 0.3333333333333333 * vz * x * z) - + 2.178553132241672 - * f[..., 177] - * (vx * vy ^ 2 * x * z - 0.3333333333333333 * vx * x * z) - + 1.257788237343632 * f[..., 118] * (vy ^ 2 * x * z - 0.3333333333333333 * x * z) - + 2.178553132241672 - * f[..., 173] - * (vx ^ 2 * vy * x * z - 0.3333333333333333 * vy * x * z) - + 1.257788237343632 * f[..., 103] * (vx ^ 2 * x * z - 0.3333333333333333 * x * z) - + 2.178553132241672 - * f[..., 218] - * (vx * vy * vz ^ 2 * z - 0.3333333333333333 * vx * vy * z) - + 1.257788237343632 - * f[..., 151] - * (vy * vz ^ 2 * z - 0.3333333333333333 * vy * z) - + 1.257788237343632 - * f[..., 148] - * (vx * vz ^ 2 * z - 0.3333333333333333 * vx * z) - + 0.7261843774138907 * f[..., 75] * (vz ^ 2 * z - 0.3333333333333333 * z) - + 2.178553132241672 - * f[..., 208] - * (vx * vy ^ 2 * vz * z - 0.3333333333333333 * vx * vz * z) - + 1.257788237343632 - * f[..., 141] - * (vy ^ 2 * vz * z - 0.3333333333333333 * vz * z) - + 2.178553132241672 - * f[..., 202] - * (vx ^ 2 * vy * vz * z - 0.3333333333333333 * vy * vz * z) - + 1.257788237343632 - * f[..., 134] - * (vx ^ 2 * vz * z - 0.3333333333333333 * vz * z) - + 1.257788237343632 - * f[..., 122] - * (vx * vy ^ 2 * z - 0.3333333333333333 * vx * z) - + 0.7261843774138907 * f[..., 66] * (vy ^ 2 * z - 0.3333333333333333 * z) - + 1.257788237343632 - * f[..., 116] - * (vx ^ 2 * vy * z - 0.3333333333333333 * vy * z) - + 0.7261843774138907 * f[..., 59] * (vx ^ 2 * z - 0.3333333333333333 * z) - + 3.375 * f[..., 219] * vx * vy * vz * x * y * z - + 1.948557158514986 * f[..., 155] * vy * vz * x * y * z - + 1.948557158514986 * f[..., 154] * vx * vz * x * y * z - + 1.125 * f[..., 83] * vz * x * y * z - + 1.948557158514986 * f[..., 153] * vx * vy * x * y * z - + 1.125 * f[..., 79] * vy * x * y * z - + 1.125 * f[..., 78] * vx * x * y * z - + 0.6495190528383289 * f[..., 28] * x * y * z - + 1.948557158514986 * f[..., 158] * vx * vy * vz * y * z - + 1.125 * f[..., 89] * vy * vz * y * z - + 1.125 * f[..., 86] * vx * vz * y * z - + 0.6495190528383289 * f[..., 40] * vz * y * z - + 1.125 * f[..., 82] * vx * vy * y * z - + 0.6495190528383289 * f[..., 34] * vy * y * z - + 0.6495190528383289 * f[..., 31] * vx * y * z - + 0.375 * f[..., 9] * y * z - + 1.948557158514986 * f[..., 157] * vx * vy * vz * x * z - + 1.125 * f[..., 88] * vy * vz * x * z - + 1.125 * f[..., 85] * vx * vz * x * z - + 0.6495190528383289 * f[..., 39] * vz * x * z - + 1.125 * f[..., 81] * vx * vy * x * z - + 0.6495190528383289 * f[..., 33] * vy * x * z - + 0.6495190528383289 * f[..., 30] * vx * x * z - + 0.375 * f[..., 8] * x * z - + 1.125 * f[..., 92] * vx * vy * vz * z - + 0.6495190528383289 * f[..., 46] * vy * vz * z - + 0.6495190528383289 * f[..., 43] * vx * vz * z - + 0.375 * f[..., 19] * vz * z - + 0.6495190528383289 * f[..., 37] * vx * vy * z - + 0.375 * f[..., 15] * vy * z - + 0.375 * f[..., 12] * vx * z - + 0.2165063509461096 * f[..., 3] * z - + 3.773364712030896 - * f[..., 233] - * (vx * vy * vz * x * y ^ 2 - 0.3333333333333333 * vx * vy * vz * x) - + 2.178553132241672 - * f[..., 192] - * (vy * vz * x * y ^ 2 - 0.3333333333333333 * vy * vz * x) - + 2.178553132241672 - * f[..., 183] - * (vx * vz * x * y ^ 2 - 0.3333333333333333 * vx * vz * x) - + 1.257788237343632 * f[..., 124] * (vz * x * y ^ 2 - 0.3333333333333333 * vz * x) - + 2.178553132241672 - * f[..., 167] - * (vx * vy * x * y ^ 2 - 0.3333333333333333 * vx * vy * x) - + 1.257788237343632 * f[..., 106] * (vy * x * y ^ 2 - 0.3333333333333333 * vy * x) - + 1.257788237343632 * f[..., 97] * (vx * x * y ^ 2 - 0.3333333333333333 * vx * x) - + 0.7261843774138907 * f[..., 49] * (x * y ^ 2 - 0.3333333333333333 * x) - + 2.178553132241672 - * f[..., 198] - * (vx * vy * vz * y ^ 2 - 0.3333333333333333 * vx * vy * vz) - + 1.257788237343632 - * f[..., 136] - * (vy * vz * y ^ 2 - 0.3333333333333333 * vy * vz) - + 1.257788237343632 - * f[..., 130] - * (vx * vz * y ^ 2 - 0.3333333333333333 * vx * vz) - + 0.7261843774138907 * f[..., 69] * (vz * y ^ 2 - 0.3333333333333333 * vz) - + 1.257788237343632 - * f[..., 112] - * (vx * vy * y ^ 2 - 0.3333333333333333 * vx * vy) - + 0.7261843774138907 * f[..., 61] * (vy * y ^ 2 - 0.3333333333333333 * vy) - + 0.7261843774138907 * f[..., 55] * (vx * y ^ 2 - 0.3333333333333333 * vx) - + 0.4192627457812106 * f[..., 23] * (y ^ 2 - 0.3333333333333333) - + 3.773364712030896 - * f[..., 232] - * (vx * vy * vz * x ^ 2 * y - 0.3333333333333333 * vx * vy * vz * y) - + 2.178553132241672 - * f[..., 191] - * (vy * vz * x ^ 2 * y - 0.3333333333333333 * vy * vz * y) - + 2.178553132241672 - * f[..., 182] - * (vx * vz * x ^ 2 * y - 0.3333333333333333 * vx * vz * y) - + 1.257788237343632 * f[..., 123] * (vz * x ^ 2 * y - 0.3333333333333333 * vz * y) - + 2.178553132241672 - * f[..., 166] - * (vx * vy * x ^ 2 * y - 0.3333333333333333 * vx * vy * y) - + 1.257788237343632 * f[..., 105] * (vy * x ^ 2 * y - 0.3333333333333333 * vy * y) - + 1.257788237343632 * f[..., 96] * (vx * x ^ 2 * y - 0.3333333333333333 * vx * y) - + 0.7261843774138907 * f[..., 48] * (x ^ 2 * y - 0.3333333333333333 * y) - + 3.773364712030896 - * f[..., 247] - * (vx * vy * vz ^ 2 * x * y - 0.3333333333333333 * vx * vy * x * y) - + 2.178553132241672 - * f[..., 213] - * (vy * vz ^ 2 * x * y - 0.3333333333333333 * vy * x * y) - + 2.178553132241672 - * f[..., 210] - * (vx * vz ^ 2 * x * y - 0.3333333333333333 * vx * x * y) - + 1.257788237343632 * f[..., 143] * (vz ^ 2 * x * y - 0.3333333333333333 * x * y) - + 3.773364712030896 - * f[..., 242] - * (vx * vy ^ 2 * vz * x * y - 0.3333333333333333 * vx * vz * x * y) - + 2.178553132241672 - * f[..., 203] - * (vy ^ 2 * vz * x * y - 0.3333333333333333 * vz * x * y) - + 3.773364712030896 - * f[..., 238] - * (vx ^ 2 * vy * vz * x * y - 0.3333333333333333 * vy * vz * x * y) - + 2.178553132241672 - * f[..., 188] - * (vx ^ 2 * vz * x * y - 0.3333333333333333 * vz * x * y) - + 2.178553132241672 - * f[..., 176] - * (vx * vy ^ 2 * x * y - 0.3333333333333333 * vx * x * y) - + 1.257788237343632 * f[..., 117] * (vy ^ 2 * x * y - 0.3333333333333333 * x * y) - + 2.178553132241672 - * f[..., 172] - * (vx ^ 2 * vy * x * y - 0.3333333333333333 * vy * x * y) - + 1.257788237343632 * f[..., 102] * (vx ^ 2 * x * y - 0.3333333333333333 * x * y) - + 2.178553132241672 - * f[..., 217] - * (vx * vy * vz ^ 2 * y - 0.3333333333333333 * vx * vy * y) - + 1.257788237343632 - * f[..., 150] - * (vy * vz ^ 2 * y - 0.3333333333333333 * vy * y) - + 1.257788237343632 - * f[..., 147] - * (vx * vz ^ 2 * y - 0.3333333333333333 * vx * y) - + 0.7261843774138907 * f[..., 74] * (vz ^ 2 * y - 0.3333333333333333 * y) - + 2.178553132241672 - * f[..., 207] - * (vx * vy ^ 2 * vz * y - 0.3333333333333333 * vx * vz * y) - + 1.257788237343632 - * f[..., 140] - * (vy ^ 2 * vz * y - 0.3333333333333333 * vz * y) - + 2.178553132241672 - * f[..., 201] - * (vx ^ 2 * vy * vz * y - 0.3333333333333333 * vy * vz * y) - + 1.257788237343632 - * f[..., 133] - * (vx ^ 2 * vz * y - 0.3333333333333333 * vz * y) - + 1.257788237343632 - * f[..., 121] - * (vx * vy ^ 2 * y - 0.3333333333333333 * vx * y) - + 0.7261843774138907 * f[..., 65] * (vy ^ 2 * y - 0.3333333333333333 * y) - + 1.257788237343632 - * f[..., 115] - * (vx ^ 2 * vy * y - 0.3333333333333333 * vy * y) - + 0.7261843774138907 * f[..., 58] * (vx ^ 2 * y - 0.3333333333333333 * y) - + 1.948557158514986 * f[..., 156] * vx * vy * vz * x * y - + 1.125 * f[..., 87] * vy * vz * x * y - + 1.125 * f[..., 84] * vx * vz * x * y - + 0.6495190528383289 * f[..., 38] * vz * x * y - + 1.125 * f[..., 80] * vx * vy * x * y - + 0.6495190528383289 * f[..., 32] * vy * x * y - + 0.6495190528383289 * f[..., 29] * vx * x * y - + 0.375 * f[..., 7] * x * y - + 1.125 * f[..., 91] * vx * vy * vz * y - + 0.6495190528383289 * f[..., 45] * vy * vz * y - + 0.6495190528383289 * f[..., 42] * vx * vz * y - + 0.375 * f[..., 18] * vz * y - + 0.6495190528383289 * f[..., 36] * vx * vy * y - + 0.375 * f[..., 14] * vy * y - + 0.375 * f[..., 11] * vx * y - + 0.2165063509461096 * f[..., 2] * y - + 2.178553132241672 - * f[..., 197] - * (vx * vy * vz * x ^ 2 - 0.3333333333333333 * vx * vy * vz) - + 1.257788237343632 - * f[..., 135] - * (vy * vz * x ^ 2 - 0.3333333333333333 * vy * vz) - + 1.257788237343632 - * f[..., 129] - * (vx * vz * x ^ 2 - 0.3333333333333333 * vx * vz) - + 0.7261843774138907 * f[..., 68] * (vz * x ^ 2 - 0.3333333333333333 * vz) - + 1.257788237343632 - * f[..., 111] - * (vx * vy * x ^ 2 - 0.3333333333333333 * vx * vy) - + 0.7261843774138907 * f[..., 60] * (vy * x ^ 2 - 0.3333333333333333 * vy) - + 0.7261843774138907 * f[..., 54] * (vx * x ^ 2 - 0.3333333333333333 * vx) - + 0.4192627457812106 * f[..., 22] * (x ^ 2 - 0.3333333333333333) - + 2.178553132241672 - * f[..., 216] - * (vx * vy * vz ^ 2 * x - 0.3333333333333333 * vx * vy * x) - + 1.257788237343632 - * f[..., 149] - * (vy * vz ^ 2 * x - 0.3333333333333333 * vy * x) - + 1.257788237343632 - * f[..., 146] - * (vx * vz ^ 2 * x - 0.3333333333333333 * vx * x) - + 0.7261843774138907 * f[..., 73] * (vz ^ 2 * x - 0.3333333333333333 * x) - + 2.178553132241672 - * f[..., 206] - * (vx * vy ^ 2 * vz * x - 0.3333333333333333 * vx * vz * x) - + 1.257788237343632 - * f[..., 139] - * (vy ^ 2 * vz * x - 0.3333333333333333 * vz * x) - + 2.178553132241672 - * f[..., 200] - * (vx ^ 2 * vy * vz * x - 0.3333333333333333 * vy * vz * x) - + 1.257788237343632 - * f[..., 132] - * (vx ^ 2 * vz * x - 0.3333333333333333 * vz * x) - + 1.257788237343632 - * f[..., 120] - * (vx * vy ^ 2 * x - 0.3333333333333333 * vx * x) - + 0.7261843774138907 * f[..., 64] * (vy ^ 2 * x - 0.3333333333333333 * x) - + 1.257788237343632 - * f[..., 114] - * (vx ^ 2 * vy * x - 0.3333333333333333 * vy * x) - + 0.7261843774138907 * f[..., 57] * (vx ^ 2 * x - 0.3333333333333333 * x) - + 1.125 * f[..., 90] * vx * vy * vz * x - + 0.6495190528383289 * f[..., 44] * vy * vz * x - + 0.6495190528383289 * f[..., 41] * vx * vz * x - + 0.375 * f[..., 17] * vz * x - + 0.6495190528383289 * f[..., 35] * vx * vy * x - + 0.375 * f[..., 13] * vy * x - + 0.375 * f[..., 10] * vx * x - + 0.2165063509461096 * f[..., 1] * x - + 1.257788237343632 - * f[..., 152] - * (vx * vy * vz ^ 2 - 0.3333333333333333 * vx * vy) - + 0.7261843774138907 * f[..., 77] * (vy * vz ^ 2 - 0.3333333333333333 * vy) - + 0.7261843774138907 * f[..., 76] * (vx * vz ^ 2 - 0.3333333333333333 * vx) - + 0.4192627457812106 * f[..., 27] * (vz ^ 2 - 0.3333333333333333) - + 1.257788237343632 - * f[..., 142] - * (vx * vy ^ 2 * vz - 0.3333333333333333 * vx * vz) - + 0.7261843774138907 * f[..., 72] * (vy ^ 2 * vz - 0.3333333333333333 * vz) - + 1.257788237343632 - * f[..., 138] - * (vx ^ 2 * vy * vz - 0.3333333333333333 * vy * vz) - + 0.7261843774138907 * f[..., 71] * (vx ^ 2 * vz - 0.3333333333333333 * vz) - + 0.6495190528383289 * f[..., 47] * vx * vy * vz - + 0.375 * f[..., 21] * vy * vz - + 0.375 * f[..., 20] * vx * vz - + 0.2165063509461096 * f[..., 6] * vz - + 0.7261843774138907 * f[..., 67] * (vx * vy ^ 2 - 0.3333333333333333 * vx) - + 0.4192627457812106 * f[..., 26] * (vy ^ 2 - 0.3333333333333333) - + 0.7261843774138907 * f[..., 63] * (vx ^ 2 * vy - 0.3333333333333333 * vy) - + 0.375 * f[..., 16] * vx * vy - + 0.2165063509461096 * f[..., 5] * vy - + 0.4192627457812106 * f[..., 25] * (vx ^ 2 - 0.3333333333333333) - + 0.2165063509461096 * f[..., 4] * vx - + 0.125 * f[..., 0] - ) - - -# end - - -def _expand_6d3p(f, x, y, z, vx, vy, vz): - return - - -# end - - -def _expand_6d4p(f, x, y, z, vx, vy, vz): - return - - -# end - -expand_6d = [_expand_6d1p, _expand_6d2p, _expand_6d3p, _expand_6d4p] diff --git a/src/postgkyl/numerics/__init__.py b/src/postgkyl/numerics/__init__.py new file mode 100644 index 00000000..12832d76 --- /dev/null +++ b/src/postgkyl/numerics/__init__.py @@ -0,0 +1,75 @@ +"""Pure NumPy/SciPy helpers -- no internal imports (the leaf-most layer).""" + +from .idx_parser import idx_parser +from .elementwise import grids_compatible, grid_is_prefix +from .calculus import integrate +from .mag_sq import mag_sq +from .rel_change import rel_change +from .rotation_matrix import rotation_matrix +from .fft import fft, init_polar, polar_isotropic +from .fit import ( + FIT_FUNCTIONS, + FIT_NDIM, + RPN_OPERATORS, + RPN_FUNCTIONS, + linear, + quadratic, + plane, + quadratic2d, + exp_plateau, + gaussian, + power, + sinusoid, + tanh_transition, + exp2, + rpn_param_names, + rpn_ndim, + fit_evaluate, + fit, + auto_guess, + fit_best_window, +) +from .filters import fft_filtering, butter_filtering +from .ev_ops import cmds as ev_cmds +from .grid_centering import nodal_to_cell_centered_grid +from .downsample import downsample +from .natural_sort import natural_sort_key + +__all__ = [ + "idx_parser", + "grids_compatible", + "grid_is_prefix", + "integrate", + "mag_sq", + "rel_change", + "rotation_matrix", + "fft", + "init_polar", + "polar_isotropic", + "FIT_FUNCTIONS", + "FIT_NDIM", + "RPN_OPERATORS", + "RPN_FUNCTIONS", + "linear", + "quadratic", + "plane", + "quadratic2d", + "exp_plateau", + "gaussian", + "power", + "sinusoid", + "tanh_transition", + "exp2", + "rpn_param_names", + "rpn_ndim", + "fit_evaluate", + "fit", + "auto_guess", + "fit_best_window", + "fft_filtering", + "butter_filtering", + "ev_cmds", + "nodal_to_cell_centered_grid", + "downsample", + "natural_sort_key", +] diff --git a/src/postgkyl/numerics/calculus.py b/src/postgkyl/numerics/calculus.py new file mode 100644 index 00000000..9518a0fe --- /dev/null +++ b/src/postgkyl/numerics/calculus.py @@ -0,0 +1,98 @@ +"""Trapezoidal-style integration over a nodal grid (pure NumPy). + +``grad``/``div``/``curl`` are deliberately absent: the ``src_bak`` originals +are unimplemented placeholders (``...`` bodies, no arguments) -- there is no +real numerics to port. The vector-calculus operators that *are* implemented +live in :mod:`postgkyl.numerics.ev_ops` (``divergence``/``curl``/``grad``), +expressed the same way, over ``(grid, values)`` pairs. +""" + +from __future__ import annotations + +import numpy as np + + +def _split_axis_string(axis: str) -> tuple: + """Parse a comma-separated (``"0,1"``) or colon-sliced (``"0:2"``) axis + string, or a bare integer string, into a tuple of integer axes. + + Shared with :func:`postgkyl.numerics.ev_ops._parse_axis`, whose outer + type-dispatch differs (it also accepts ``float``/``np.ndarray``/``"all"``) + but delegates this exact string-parsing branch here, so the comma/colon + grammar has one home (Doctrine V) instead of two copies that could drift. + """ + if len(axis.split(",")) > 1: + return tuple(int(a) for a in axis.split(",")) + if len(axis.split(":")) == 2: + lo, hi = axis.split(":") + return tuple(range(int(lo), int(hi))) + return (int(axis), ) + + +def parse_axis(axis: int | tuple | str | None, num_dims: int) -> tuple: + """Turn an axis selector into a tuple of integer axes.""" + if axis is None: + return tuple(range(num_dims)) + if isinstance(axis, int): + return (axis, ) + if isinstance(axis, tuple): + return axis + if isinstance(axis, str): + return _split_axis_string(axis) + raise TypeError( + "'axis' needs to be integer, tuple, string of comma separated " + "integers, or a slice ('int:int')") + + +def integrate( + grid: list[np.ndarray], + values: np.ndarray, + axis: int | tuple | str | None = None +) -> tuple[list[np.ndarray], np.ndarray]: + """Integrate cell-centered-average data over one or more axes. + + Uses the NumPy dot product against the cell widths (trapezoidal for + nodal/edge grids, exact for cell-centered-average data); works for + nonuniform meshes. True DG integration is not implemented here -- this + mirrors the legacy behaviour exactly. + + Args: + grid: Nodal (edge) coordinate arrays, one per spatial dimension. + values: Data array; the last axis is components, the rest are spatial. + axis: Axis (or axes) to integrate over: an ``int``, a ``tuple`` of + ``int``, a comma-separated string (``"0,1"``), a colon slice string + (``"0:2"``), or ``None`` (integrate over every spatial axis). + + Returns: + ``(grid, values)`` with the integrated axes collapsed to a single, + grid-mean cell and ``values`` reduced accordingly (shape retained via + ``expand_dims``). + + Raises: + TypeError: If ``axis`` is not an int, tuple, or string. + """ + grid = list(grid) + values = np.copy(values) + axis = parse_axis(axis, len(grid)) + + # Get dz elements + dz = [] + for d, coord in enumerate(grid): + dz.append(coord[1:] - coord[:-1]) + if len(coord) > 1 and len(coord) == values.shape[d]: + dz[-1] = np.append(dz[-1], dz[-1][-1]) + + # Integration assuming values are cell centered averages + # Should work for nonuniform meshes + for ax in sorted(axis, reverse=True): + if len(grid[ax]) > 1: + values = np.moveaxis(values, ax, -1) + values = np.dot(values, dz[ax]) + else: + values = values.mean(axis=ax) + + for ax in sorted(axis): + grid[ax] = np.array([grid[ax].mean()]) + values = np.expand_dims(values, ax) + + return grid, values diff --git a/src/postgkyl/numerics/curvilinear.py b/src/postgkyl/numerics/curvilinear.py new file mode 100644 index 00000000..049b15a5 --- /dev/null +++ b/src/postgkyl/numerics/curvilinear.py @@ -0,0 +1,123 @@ +"""Pure NumPy geometry for curvilinear (non-separable, ``.map()``-produced) +grid blocks. + +A curvilinear block (see ``postgkyl.operations.map``, ``space="conf"`` with +``m > 1``) stores ``m`` physical-coordinate arrays, each shaped like the +block's own joint ``m``-D nodal (edge) grid -- unlike a separable axis, no +single 1-D coordinate array exists per dimension, so a plain ``np.gradient`` +or coordinate-difference has no meaning. These helpers compute the local +geometry needed to differentiate or integrate data on such a block via the +chain rule / change of variables instead: the index-space Jacobian +(``jacobian``), its determinant as the physical cell volume (``cell_volume``, +used by ``integrate``), and the resulting physical-space gradient +(``physical_gradient``, used by ``differentiate``). + +The Jacobian is evaluated with unit index spacing (one cell = one index +step) rather than the block's original computational grid spacing, which +``.map()`` does not retain. This is exact for ``cell_volume``/ +``physical_gradient``'s purposes: both only ever use the Jacobian in a ratio +(inverted against a same-parametrization gradient, or as a *relative* cell +weight normalized by the sum of the block's cells), so the arbitrary choice +of index units cancels out. +""" + +from __future__ import annotations + +import numpy as np + + +def cell_center(nodal: np.ndarray) -> np.ndarray: + """Average an ``ndim``-D nodal (edge) array over its ``2**ndim`` corners. + + The curvilinear analogue of the ``0.5 * (coord[1:] + coord[:-1])`` + cell-center convention used elsewhere for a single (separable) axis: + reduces every axis' length by one. + """ + out = nodal + for ax in range(out.ndim): + lo = tuple( + slice(0, -1) if k == ax else slice(None) for k in range(out.ndim)) + hi = tuple( + slice(1, None) if k == ax else slice(None) for k in range(out.ndim)) + out = 0.5 * (out[lo] + out[hi]) + return out + + +def jacobian(block_coords: list) -> np.ndarray: + """The index-space Jacobian of an ``m``-D curvilinear block. + + Args: + block_coords: the block's ``m`` physical-coordinate arrays (one per + mapped dimension, in the block's own local-axis order), each of the + block's own nodal (edge) shape. + + Returns: + ``J`` of shape ``cells_shape + (m, m)``, where ``cells_shape`` is + ``block_coords[0]``'s shape reduced by one per axis (cell-centered) and + ``J[..., i, j] = d(block_coords[i]) / d(local cell index j)``, evaluated + by central differences at unit index spacing. + """ + m = len(block_coords) + centers = [cell_center(c) for c in block_coords] + shape = centers[0].shape + J = np.empty(shape + (m, m)) + for i in range(m): + for j in range(m): + J[..., i, j] = np.gradient(centers[i], axis=j, edge_order=2) + return J + + +def cell_volume(block_coords: list) -> np.ndarray: + """Per-cell physical volume (area in 2-D) of an ``m``-D curvilinear block. + + The change-of-variables volume element (the Jacobian determinant): exact + for a bilinear/trilinear cell, second-order accurate otherwise -- matching + this codebase's numerical (not exact) differentiate/integrate philosophy. + Unit index spacing already gives the *physical* cell volume directly + (not merely a relative one up to some missing scale): ``block_coords`` + holds true physical coordinates against a unit-index abscissa, so a + central difference of one index step is ``dxi`` times the continuous + ``d(physical)/d(index)`` derivative, and that same ``dxi`` factor appears + once per row of the Jacobian -- i.e. ``m`` times in its determinant, + exactly cancelling the ``m``-fold ``1/dxi`` of the physical volume + element ``dxi_0 * dxi_1 * ... `` No separate pre-map cell-width metadata + is needed (or, unlike ``physical_gradient``, would even help: here it + would double-count the very same factor). + + Shape: the block's ``cells_shape`` (one entry per mapped dimension). + """ + return np.abs(np.linalg.det(jacobian(block_coords))) + + +def physical_gradient(block_coords: list, values: np.ndarray, + block_axes: tuple) -> np.ndarray: + """The physical-space gradient of ``values`` along a curvilinear block's + directions, via the chain rule ``grad_x f = (J^-1)^T grad_xi f``. + + Args: + block_coords: the block's ``m`` physical-coordinate arrays, each of the + block's own nodal (edge) shape, in ``block_axes`` order. + values: the dataset's cell-centered values (any number of axes; the + block's cells occupy the absolute axes named in ``block_axes``). + block_axes: the absolute axis of ``values`` differentiated by each of + ``block_coords``'s local dimensions, in order. + + Returns: + ``values.shape + (m,)``: the physical derivative along each of the + block's ``m`` directions (``block_axes`` order) on a new trailing axis; + every other axis keeps ``values``'s own shape and position. + """ + m = len(block_coords) + jinv = np.linalg.inv(jacobian(block_coords)) # cells_shape + (m, m) + + moved = np.moveaxis(values, block_axes, range(m)) + dfdxi = np.stack([np.gradient(moved, axis=j, edge_order=2) for j in range(m)], + axis=-1) + + # jinv only varies over the block's own m axes; insert size-1 axes for + # every other axis of `moved` (now trailing, after the moveaxis above) so + # it broadcasts against dfdxi positionally. + n_between = moved.ndim - m + jinv = jinv.reshape(jinv.shape[:m] + (1, ) * n_between + (m, m)) + out_moved = np.einsum("...ji,...j->...i", jinv, dfdxi) + return np.moveaxis(out_moved, range(m), block_axes) diff --git a/src/postgkyl/numerics/downsample.py b/src/postgkyl/numerics/downsample.py new file mode 100644 index 00000000..11db53a7 --- /dev/null +++ b/src/postgkyl/numerics/downsample.py @@ -0,0 +1,67 @@ +"""Downsample same-shape arrays so no axis exceeds a configured maximum.""" + +from __future__ import annotations + +import numpy as np + + +def downsample(*arrays: np.ndarray, + maximum_points_per_axis: int = 0) -> tuple[np.ndarray, ...]: + """Downsample same-shape arrays so no axis exceeds ``maximum_points_per_axis``. + + Dimension-agnostic: works for any array dimensionality. If the arrays' + shapes disagree, or no downsampling is needed/requested, the arrays are + returned unchanged. + + Args: + *arrays: One or more arrays to downsample. All arrays must have the + same shape. + maximum_points_per_axis: The maximum number of points allowed along + any axis after downsampling. If ``0`` or negative, no downsampling + is performed. + + Returns: + A tuple of downsampled arrays corresponding to the input arrays. + + Example: + >>> x = np.linspace(0, 10, 100) + >>> value = np.random.rand(100) + >>> x_ds, value_ds = downsample(x, value, maximum_points_per_axis=20) + """ + if not arrays: + return () + + reference = arrays[0] + if maximum_points_per_axis is None or maximum_points_per_axis <= 0: + return arrays + + if reference.ndim == 0: + return arrays + + if any(arr.shape != reference.shape for arr in arrays): + return arrays + + steps = [ + max(1, int(np.ceil(size / maximum_points_per_axis))) + for size in reference.shape + ] + if max(steps) == 1: + return arrays + + def _axis_indices(size: int, step: int) -> np.ndarray: + idx = np.arange(0, size, step, dtype=int) + if idx[-1] != size - 1: + idx = np.append(idx, size - 1) + return idx + + axis_indices = [ + _axis_indices(size, step) for size, step in zip(reference.shape, steps) + ] + + def _take_indices(arr: np.ndarray) -> np.ndarray: + out = arr + for axis, idx in enumerate(axis_indices): + out = np.take(out, idx, axis=axis) + return out + + return tuple(_take_indices(arr) for arr in arrays) diff --git a/src/postgkyl/numerics/elementwise.py b/src/postgkyl/numerics/elementwise.py new file mode 100644 index 00000000..4243c2c8 --- /dev/null +++ b/src/postgkyl/numerics/elementwise.py @@ -0,0 +1,24 @@ +"""Pure-array helpers for element-wise dataset arithmetic.""" + +from __future__ import annotations + +import numpy as np + + +def grids_compatible(grid_a: list, grid_b: list, rtol: float = 1e-9) -> bool: + """Whether two nodal grids describe the same mesh (same shapes & nodes).""" + if len(grid_a) != len(grid_b): + return False + return all(a.shape == b.shape and np.allclose(a, b, rtol=rtol) + for a, b in zip(grid_a, grid_b)) + + +def grid_is_prefix(small: list, big: list, rtol: float = 1e-9) -> bool: + """Whether ``small`` is exactly the leading dimensions of ``big`` (same + shapes & nodes) -- the conf-space/phase-space compatibility check for + cross-basis (conf x phase) operations, where a phase-space grid extends a + lower-dimensional conf-space grid with extra (velocity-space) dimensions.""" + if not 0 < len(small) < len(big): + return False + return all(a.shape == b.shape and np.allclose(a, b, rtol=rtol) + for a, b in zip(small, big[:len(small)])) diff --git a/src/postgkyl/numerics/ev_ops.py b/src/postgkyl/numerics/ev_ops.py new file mode 100644 index 00000000..b8d7ca73 --- /dev/null +++ b/src/postgkyl/numerics/ev_ops.py @@ -0,0 +1,552 @@ +"""RPN operator registry for the ``ev`` verb (pure ``(grid, values)`` functions). + +This is the numeric core behind the ``ev`` expression evaluator. Each +operator is a pure function ``f(in_grid, in_values) -> ([out_grid], [out_values])`` +over plain Python lists / NumPy arrays -- no ``GData`` dependency. The +``cmds`` table maps each RPN token to its arity (``num_in``/``num_out``) +and function; the stack machine that drives them lives in the ``operations`` +layer's ``ev`` verb (layer 07), which can consume this table unchanged. + +Every operator here is expressible over plain arrays; none needed a +``NotImplementedError`` GData-only placeholder. +""" + +from __future__ import annotations + +import numpy as np + +from .calculus import _split_axis_string +from .idx_parser import idx_parser + + +def _get_grid(grid0, grid1): + if grid0 is not None and grid1 is not None: + return grid0 if len(grid0) > len(grid1) else grid1 + if grid0 is not None: + return grid0 + if grid1 is not None: + return grid1 + return None + + +def add(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + out_values = in_values[0] + in_values[1] + return [out_grid], [out_values] + + +def subtract(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + out_values = in_values[1] - in_values[0] + return [out_grid], [out_values] + + +def mult(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + a, b = in_values[1], in_values[0] + if np.array_equal(a.shape, b.shape) or len(a.shape) == 0 or len(b.shape) == 0: + out_values = a * b + else: + # When multiplying a phase-space and a conf-space field, the + # dimensions do not match. NumPy broadcasting requires the *trailing* + # indices to match, which is the opposite of what we have here (the + # *leading* indices match) -- so transpose, multiply, transpose back. + out_values = (a.transpose() * b.transpose()).transpose() + return [out_grid], [out_values] + + +def dot(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + out_values = np.sum(in_values[1] * in_values[0], axis=-1)[..., np.newaxis] + return [out_grid], [out_values] + + +def divide(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + a, b = in_values[1], in_values[0] + if np.array_equal(a.shape, b.shape) or len(a.shape) == 0 or len(b.shape) == 0: + out_values = a / b + else: + # See the 'mult' comment above. + out_values = (a.transpose() / b.transpose()).transpose() + return [out_grid], [out_values] + + +def sqrt(in_grid, in_values): + return [in_grid[0]], [np.sqrt(in_values[0])] + + +def psin(in_grid, in_values): + return [in_grid[0]], [np.sin(in_values[0])] + + +def pcos(in_grid, in_values): + return [in_grid[0]], [np.cos(in_values[0])] + + +def ptan(in_grid, in_values): + return [in_grid[0]], [np.tan(in_values[0])] + + +def absolute(in_grid, in_values): + return [in_grid[0]], [np.abs(in_values[0])] + + +def log(in_grid, in_values): + return [in_grid[0]], [np.log(in_values[0])] + + +def log10(in_grid, in_values): + return [in_grid[0]], [np.log10(in_values[0])] + + +def minimum(in_grid, in_values): + out_values = np.atleast_1d(np.nanmin(in_values[0])) + return [[]], [out_values] + + +def minimum2(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + out_values = np.fmin(in_values[0], in_values[1]) + return [out_grid], [out_values] + + +def maximum(in_grid, in_values): + out_values = np.atleast_1d(np.nanmax(in_values[0])) + return [[]], [out_values] + + +def maximum2(in_grid, in_values): + out_grid = _get_grid(in_grid[0], in_grid[1]) + out_values = np.fmax(in_values[0], in_values[1]) + return [out_grid], [out_values] + + +def mean(in_grid, in_values): + out_values = np.atleast_1d(np.mean(in_values[0])) + return [[]], [out_values] + + +def power(in_grid, in_values): + out_grid = in_grid[1] + out_values = np.power(in_values[1], in_values[0]) + return [out_grid], [out_values] + + +def sq(in_grid, in_values): + return [in_grid[0]], [in_values[0]**2] + + +def exp(in_grid, in_values): + return [in_grid[0]], [np.exp(in_values[0])] + + +def length(in_grid, in_values): + ax = int(in_values[0]) + ln = in_grid[1][ax][-1] - in_grid[1][ax][0] + if len(in_grid[1][ax]) == in_values[1].shape[ax]: + ln += in_grid[1][ax][1] - in_grid[1][ax][0] + return [[]], [ln] + + +def grad(in_grid, in_values): + out_grid = in_grid[0] + nd = len(in_values[0].shape) - 1 + out_shape = list(in_values[0].shape) + nc = in_values[0].shape[-1] + out_shape[-1] = nc * nd + out_values = np.zeros(out_shape) + + for d in range(nd): + zc = 0.5 * (in_grid[0][d][1:] + in_grid[0][d][:-1]) # cell centered values + out_values[..., d * nc:(d + 1) * nc] = np.gradient(in_values[0], + zc, + edge_order=2, + axis=d) + return [out_grid], [out_values] + + +def grad2(in_grid, in_values): + out_grid = in_grid[1] + ax = in_values[0] + if isinstance(ax, str) and ":" in ax: + lo, up = ax.split(":") + rng = range(int(lo), int(up)) + elif isinstance(ax, str): + rng = tuple(int(i) for i in ax.split(",")) + else: + rng = range(int(ax), int(ax + 1)) + + num_dims = len(rng) + out_shape = list(in_values[1].shape) + num_comps = in_values[1].shape[-1] + out_shape[-1] = out_shape[-1] * num_dims + out_values = np.zeros(out_shape) + + for cnt, d in enumerate(rng): + zc = 0.5 * (in_grid[1][d][1:] + in_grid[1][d][:-1]) # cell centered values + out_values[..., cnt * num_comps:(cnt + 1) * num_comps] = np.gradient( + in_values[1], zc, edge_order=2, axis=d) + return [out_grid], [out_values] + + +def _parse_axis(axis) -> tuple: + if isinstance(axis, float): + return (int(axis), ) + if isinstance(axis, tuple): + return axis + if isinstance(axis, np.ndarray): + return (int(axis), ) + if isinstance(axis, str): + if axis == "all": + return None # resolved against num_dims by the caller + return _split_axis_string(axis) + raise TypeError( + "'axis' needs to be integer, tuple, string of comma separated " + "integers, or a slice ('int:int')") + + +def integrate(in_grid, in_values, avg=False): + grid = in_grid[1].copy() + values = np.array(in_values[1]) + + axis = _parse_axis(in_values[0]) + if axis is None: + axis = tuple(range(len(grid))) + + dz = [] + for d, coord in enumerate(grid): + dz.append(coord[1:] - coord[:-1]) + if len(coord) == values.shape[d]: + dz[-1] = np.append(dz[-1], dz[-1][-1]) + + # Integration assuming values are cell centered averages + # Should work for nonuniform meshes + for ax in sorted(axis, reverse=True): + values = np.moveaxis(values, ax, -1) + values = np.dot(values, dz[ax]) + for ax in sorted(axis): + grid[ax] = np.array([0]) + values = np.expand_dims(values, ax) + if avg: + ln = in_grid[1][ax][-1] - in_grid[1][ax][0] + if len(in_grid[1][ax]) == in_values[1].shape[ax]: + ln += in_grid[1][ax][1] - in_grid[1][ax][0] + values = values / ln + return [grid], [values] + + +def average(in_grid, in_values): + return integrate(in_grid, in_values, True) + + +def divergence(in_grid, in_values): + out_grid = in_grid[0] + num_dims = len(in_grid[0]) + num_comps = in_values[0].shape[-1] + if num_comps > num_dims: + # src_bak warned and computed a partial result (using only the first + # num_dims components) here; this raises instead per PYTHON_PRINCIPLES §10. + raise ValueError( + f"ERROR in 'evaluate div': Length of the provided vector ({num_comps:d}) " + f"is longer than number of dimensions ({num_dims:d}).") + out_shape = list(in_values[0].shape) + out_shape[-1] = 1 + out_values = np.zeros(out_shape) + for d in range(num_dims): + zc = 0.5 * (in_grid[0][d][1:] + in_grid[0][d][:-1]) # cell centered values + out_values[..., 0] = out_values[..., 0] + np.gradient( + in_values[0][..., d], zc, edge_order=2, axis=d) + return [out_grid], [out_values] + + +def curl(in_grid, in_values): + out_grid = in_grid[0] + num_dims = len(in_grid[0]) + num_comps = in_values[0].shape[-1] + + out_shape = list(in_values[0].shape) + + if num_dims == 1: + if num_comps != 3: + raise ValueError( + f"ERROR in 'evaluate curl': Curl in 1D requires 3-component input and " + f"{num_comps:d}-component field was provided.") + zc0 = 0.5 * (in_grid[0][0][1:] + in_grid[0][0][:-1]) + out_values = np.zeros(out_shape) + out_values[..., 1] = -np.gradient( + in_values[0][..., 2], zc0, edge_order=2, axis=0) + out_values[..., 2] = np.gradient(in_values[0][..., 1], + zc0, + edge_order=2, + axis=0) + elif num_dims == 2: + zc0 = 0.5 * (in_grid[0][0][1:] + in_grid[0][0][:-1]) + zc1 = 0.5 * (in_grid[0][1][1:] + in_grid[0][1][:-1]) + if num_comps < 2: + raise ValueError( + f"ERROR in 'evaluate curl': Length of the provided vector ({num_comps:d}) " + f"is smaller than number of dimensions ({num_dims:d}). Curl can't " + f"be calculated.") + elif num_comps == 2: + # A 2D vector field: curl reduces to the single in-plane (z) component. + # This is the normal, expected input for 2D curl -- not an anomaly. + out_shape[-1] = 1 + out_values = np.zeros(out_shape) + out_values[..., 0] = np.gradient( + in_values[0][..., 1], zc0, edge_order=2, axis=0) - np.gradient( + in_values[0][..., 0], zc1, edge_order=2, axis=1) + else: + if num_comps > 3: + # src_bak warned and computed a partial result (using only the + # first 3 components) here; this raises instead per + # PYTHON_PRINCIPLES §10. + raise ValueError( + f"ERROR in 'evaluate curl': Length of the provided vector " + f"({num_comps:d}) is longer than number of dimensions " + f"({num_dims:d}).") + out_values = np.zeros(out_shape) + out_values[..., 0] = np.gradient(in_values[0][..., 2], + zc1, + edge_order=2, + axis=1) + out_values[..., 1] = -np.gradient( + in_values[0][..., 2], zc0, edge_order=2, axis=0) + out_values[..., 2] = np.gradient( + in_values[0][..., 1], zc0, edge_order=2, axis=0) - np.gradient( + in_values[0][..., 0], zc1, edge_order=2, axis=1) + else: # 3D + if num_comps > 3: + # src_bak warned and computed a partial result (using only the + # first 3 components) here; this raises instead per + # PYTHON_PRINCIPLES §10. + raise ValueError( + f"ERROR in 'evaluate curl': Length of the provided vector ({num_comps:d}) " + f"is longer than number of dimensions ({num_dims:d}).") + elif num_comps < 3: + raise ValueError( + f"ERROR in 'evaluate curl': Length of the provided vector ({num_comps:d}) " + f"is smaller than number of dimensions ({num_dims:d}). Curl can't " + f"be calculated.") + zc0 = 0.5 * (in_grid[0][0][1:] + in_grid[0][0][:-1]) + zc1 = 0.5 * (in_grid[0][1][1:] + in_grid[0][1][:-1]) + zc2 = 0.5 * (in_grid[0][2][1:] + in_grid[0][2][:-1]) + out_values = np.zeros(out_shape) + out_values[..., 0] = np.gradient( + in_values[0][..., 2], zc1, edge_order=2, axis=1) - np.gradient( + in_values[0][..., 1], zc2, edge_order=2, axis=2) + out_values[..., 1] = np.gradient( + in_values[0][..., 0], zc2, edge_order=2, axis=2) - np.gradient( + in_values[0][..., 2], zc0, edge_order=2, axis=0) + out_values[..., 2] = np.gradient( + in_values[0][..., 1], zc0, edge_order=2, axis=0) - np.gradient( + in_values[0][..., 0], zc1, edge_order=2, axis=1) + return [out_grid], [out_values] + + +def scale_comp(in_grid, in_values): + """Scale specific components of the data. + + RPN stack order: ``f comp_spec scale_factor scale_comp`` -- usage + ``f 2:4 1000 scale_comp`` scales components 2 and 3 by 1000. + + Args: + in_values[0]: Scaling factor. + in_values[1]: Component specification (a string like ``"2:4"``, or a + number). + in_values[2]: Original data array (``f``). + """ + out_grid = in_grid[2] # grid from the original data (f) + original_data = in_values[2].copy() + comp_spec = in_values[1] + scale_factor = in_values[0] + + scale_factor = scale_factor.item() + if isinstance(comp_spec, str): + comp_idx = idx_parser(comp_spec) + elif isinstance(comp_spec, np.ndarray) and comp_spec.size == 1: + comp_idx = int(comp_spec.item()) + else: + comp_idx = int(comp_spec) + + if isinstance(comp_idx, slice): + original_data[..., comp_idx] *= scale_factor + elif isinstance(comp_idx, tuple): + for idx in comp_idx: + original_data[..., idx] *= scale_factor + else: + original_data[..., comp_idx] *= scale_factor + + return [out_grid], [original_data] + + +def scale_zi_axis(in_grid, in_values): + """Scale the ``z_i`` axis of the grid. + + RPN stack order: ``f axis scale_factor scale_zi_axis`` -- usage + ``f 0 1000 scale_zi_axis`` scales the x-axis (axis 0) by 1000. + + Args: + in_values[0]: Scaling factor. + in_values[1]: Axis direction (``0``-``5``). + in_values[2]: Original data array (``f``). + """ + out_grid = in_grid[2] # grid from the original data (f) + original_data = in_values[2].copy() + idx_scale = in_values[1].item() + scale_factor = in_values[0].item() + + # NB: mutates the referenced axis array in place (matches src_bak exactly, + # including its aliasing with the caller's original grid list). + out_grid[int(idx_scale)] *= scale_factor + + return [out_grid], [original_data] + + +cmds = { + "+": { + "num_in": 2, + "num_out": 1, + "func": add + }, + "-": { + "num_in": 2, + "num_out": 1, + "func": subtract + }, + "*": { + "num_in": 2, + "num_out": 1, + "func": mult + }, + "/": { + "num_in": 2, + "num_out": 1, + "func": divide + }, + "dot": { + "num_in": 2, + "num_out": 1, + "func": dot + }, + "sqrt": { + "num_in": 1, + "num_out": 1, + "func": sqrt + }, + "sin": { + "num_in": 1, + "num_out": 1, + "func": psin + }, + "cos": { + "num_in": 1, + "num_out": 1, + "func": pcos + }, + "tan": { + "num_in": 1, + "num_out": 1, + "func": ptan + }, + "abs": { + "num_in": 1, + "num_out": 1, + "func": absolute + }, + "avg": { + "num_in": 2, + "num_out": 1, + "func": average + }, + "log": { + "num_in": 1, + "num_out": 1, + "func": log + }, + "log10": { + "num_in": 1, + "num_out": 1, + "func": log10 + }, + "max": { + "num_in": 1, + "num_out": 1, + "func": maximum + }, + "min": { + "num_in": 1, + "num_out": 1, + "func": minimum + }, + "max2": { + "num_in": 2, + "num_out": 1, + "func": maximum2 + }, + "min2": { + "num_in": 2, + "num_out": 1, + "func": minimum2 + }, + "mean": { + "num_in": 1, + "num_out": 1, + "func": mean + }, + "len": { + "num_in": 2, + "num_out": 1, + "func": length + }, + "pow": { + "num_in": 2, + "num_out": 1, + "func": power + }, + "sq": { + "num_in": 1, + "num_out": 1, + "func": sq + }, + "exp": { + "num_in": 1, + "num_out": 1, + "func": exp + }, + "grad": { + "num_in": 1, + "num_out": 1, + "func": grad + }, + "grad2": { + "num_in": 2, + "num_out": 1, + "func": grad2 + }, + "int": { + "num_in": 2, + "num_out": 1, + "func": integrate + }, + "div": { + "num_in": 1, + "num_out": 1, + "func": divergence + }, + "curl": { + "num_in": 1, + "num_out": 1, + "func": curl + }, + "scale_comp": { + "num_in": 3, + "num_out": 1, + "func": scale_comp + }, + "scale_zi_axis": { + "num_in": 3, + "num_out": 1, + "func": scale_zi_axis + }, +} diff --git a/src/postgkyl/numerics/fft.py b/src/postgkyl/numerics/fft.py new file mode 100644 index 00000000..ec811eae --- /dev/null +++ b/src/postgkyl/numerics/fft.py @@ -0,0 +1,257 @@ +"""FFT / PSD of gridded data, plus polar (shell) isotropic binning. + +Merges the legacy ``tools/fft.py``, ``tools/init_polar.py``, and +``tools/polar_isotropic.py`` into one module: :func:`fft` is the entry +point (with ``psd``/``iso`` flags), :func:`init_polar` and +:func:`polar_isotropic` are the isotropic-binning helpers it calls for +``iso=True`` and are also useful standalone. +""" + +from __future__ import annotations + +import numpy as np +import scipy.fft + + +def fft(grid: list[np.ndarray], + values: np.ndarray, + *, + psd: bool = False, + iso: bool = False) -> tuple[list[np.ndarray], np.ndarray]: + """FFT (or power spectral density, optionally isotropic) of gridded data. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension. Axes of + length <= 2 are treated as dummy dimensions and squeezed out first. + values: Data array; the last axis is components. + psd: If ``True``, return the (one-sided) power spectral density + instead of the complex FFT. + iso: If ``True`` (requires ``psd`` and exactly 3 real spatial + dimensions), additionally shell-average the PSD over polar + (isotropic) ``k``-bins and return a 1-D isotropic spectrum. + + Returns: + ``(freq, ft_values)``: ``freq`` is a list of 1-D frequency arrays (one + per surviving spatial axis, or a single polar-``k`` axis if ``iso``), + and ``ft_values`` is the (P)FT array. + + Raises: + ValueError: If ``psd`` is requested for data that is not 1-D, 2-D, or + 3-D. + """ + grid = list(grid) + values = values + + # Remove dummy dimensions + num_dims = len(grid) + idx = [d for d in range(num_dims) if len(grid[d]) <= 2] + if idx: + for i in idx[::-1]: + grid.pop(i) + values = np.squeeze(values, tuple(idx)) + num_dims = len(grid) + num_comps = values.shape[-1] + + if num_dims == 1: + N = len(grid[0]) + dx = grid[0][1] - grid[0][0] + freq = [scipy.fft.fftfreq(N, dx)] + ft_values = np.zeros(values.shape, "complex") + for comp in np.arange(num_comps): + ft_values[..., comp] = scipy.fft.fft(values[..., comp]) + + if psd: + freq[0] = freq[0][:N // 2] + ft_values = np.abs(ft_values[:N // 2, :])**2 + return freq, ft_values + + if num_dims > 3: + # src_bak raised this same message, but only from deep inside the + # ``psd`` branch -- unreachable in practice, since the fixed-size + # ``N = np.zeros(3)`` below always raises a confusing IndexError first + # for num_dims > 3, psd or not. Raise it up front instead. + raise ValueError("Only 1D, 2D, and 3D data are currently supported.") + + N = np.zeros(3, dtype=int) + dx = np.zeros(3) + freq = [] + for i in range(num_dims): + N[i] = len(grid[i]) + dx[i] = grid[i][1] - grid[i][0] + freq.append(scipy.fft.fftfreq(N[i], dx[i])) + ft_values = np.zeros(values.shape, "complex") + for comp in np.arange(num_comps): + ft_values[..., comp] = scipy.fft.fftn(values[..., comp]) + if not psd: + return freq, ft_values + + for i in range(num_dims): + freq[i] = freq[i][:N[i] // 2] + if num_dims == 2: + ft_values = np.abs(ft_values[:N[0] // 2, :N[1] // 2, :])**2 + if iso: + freq.append(0) # dummy third index, only meaningful to init_polar below + else: # num_dims == 3 (num_dims > 3 already raised above) + ft_values = np.abs(ft_values[:N[0] // 2, :N[1] // 2, :N[2] // 2, :])**2 + + if not iso: + return freq, ft_values + + nkpolar = int(np.sqrt(np.sum(N[:]**2))) + nkx = N[0] // 2 + nky = N[1] // 2 + nkz = N[2] // 2 + kx, ky, kz = freq[0], freq[1], freq[2] + akp, nbin, polar_index, _ = init_polar(nkx, nky, nkz, kx, ky, kz, nkpolar) + fft_iso = np.zeros((nkpolar, num_comps)) + for comp in np.arange(num_comps): + fft_iso[:, comp] = polar_isotropic(nkpolar, nkx, nky, nkz, polar_index, + nbin, ft_values[..., comp], kx, ky, kz) + return [akp], fft_iso + + +def init_polar(nkx, nky, nkz, kx, ky, kz, nkpolar): + """Build a polar (k-perpendicular) binning of a Cartesian wavenumber grid. + + Constructs uniformly spaced polar bins in ``k = sqrt(kx**2 + ky**2 [+ kz**2])`` + and assigns each Cartesian wavenumber cell to a bin, for later isotropic + (shell) averaging of spectra. Works for 2D grids (set ``nkz`` and ``kz`` to + ``0``) and 3D grids. + + Args: + nkx: Number of grid points along the ``kx`` axis. + nky: Number of grid points along the ``ky`` axis. + nkz: Number of grid points along the ``kz`` axis; use ``0`` for 2D data. + kx: 1D array of ``kx`` wavenumbers; ``kx[1]`` sets the spacing ``dkx``. + ky: 1D array of ``ky`` wavenumbers; ``ky[1]`` sets the spacing ``dky``. + kz: 1D array of ``kz`` wavenumbers; ``kz[1]`` sets the spacing ``dkz``. + Use ``0`` for 2D data. + nkpolar: Number of polar (radial ``k_perp``) bins to create. If ``0``, + no binning is performed and empty outputs are returned. + + Returns: + ``(akp, nbin, polar_index, akplim)`` where ``akp`` is the array of + polar bin centers (the ``k_perp`` grid), ``nbin`` is the count of + Cartesian cells assigned to each bin, ``polar_index`` is an integer + array (shape matching the Cartesian grid) giving the bin index of each + cell, and ``akplim`` is the array of polar bin edges. + """ + # if 2D, nkz and kz = 0 + + if nkpolar == 0: + akp = [] + nbin = 0 + polar_index = [] + akplim = [] + elif nkz == 0: + nbin = np.zeros(nkpolar) # Number of kx,ky in each polar bins + polar_index = np.zeros((nkx, nky), + dtype=int) # Polar index to simplify binning + if nkx == 1 and nky == 1: + # NB: src_bak wrote this as ``nkx == 1 & nky == 1``. ``&`` binds + # tighter than ``==``, so that parsed as + # ``nkx == (1 & nky) and (1 & nky) == 1`` -- true or false by the + # *parity* of nky, not by whether nkx/nky actually equal 1. Fixed to + # the evidently intended ``and``, proven by the parity-sensitive + # test in test_numerics_fft.py. + dkp = 0 + elif nkx == 1: + dkp = ky[1] + elif nky == 1: + dkp = kx[1] + else: + dkp = max(kx[1], ky[1]) + akp = (np.linspace(1, nkpolar, nkpolar)) * dkp # Kperp grid + akplim = dkp / 2 + (np.linspace(0, nkpolar, + nkpolar + 1)) * dkp # Bin limits + # Re-written to avoid loops. Necessary for large grids. + [kxg, kyg] = np.meshgrid( + ky, kx) # Deal with meshgrid weirdness (so do not have to transpose) + kp = np.sqrt(kxg**2 + kyg**2) + pn = np.where(kp >= akplim[nkpolar]) + polar_index[pn[0], pn[1]] = nkpolar - 1 + nbin[nkpolar - 1] = nbin[nkpolar - 1] + len(pn[0]) + for ik in range(0, nkpolar): + pn = np.where((kp < akplim[ik + 1]) & (kp >= akplim[ik])) + polar_index[pn[0], pn[1]] = ik + nbin[ik] = nbin[ik] + len(pn[0]) + else: + # 3D data + nbin = np.zeros(nkpolar) + polar_index = np.zeros((nkx, nky, nkz), dtype=int) + if nkx == 1 and nky == 1 and nkz == 1: + # NB: same ``&``-vs-``==``-precedence bug as the 2D branch above, + # fixed the same way. + dkp = 0 + elif nkx == 1: + dkp = max(ky[1], kz[1]) + elif nky == 1: + dkp = max(kx[1], kz[1]) + elif nkz == 1: + dkp = max(kx[1], ky[1]) + else: + dkp = max(kx[1], ky[1], kz[1]) + akp = (np.linspace(1, nkpolar, nkpolar)) * dkp # kperp grid + akplim = dkp / 2 + (np.linspace(0, nkpolar, + nkpolar + 1)) * dkp # bin limits + # Re-written to avoid loops + [kxg, kyg, kzg] = np.meshgrid(ky, kx, kz) + kp = np.sqrt(kxg**2 + kyg**2 + kzg**2) + pn = np.where(kp >= akplim[nkpolar]) + polar_index[pn[0], pn[1], pn[2]] = nkpolar - 1 + nbin[nkpolar - 1] = nbin[nkpolar - 1] + len(pn[0]) + for ik in range(0, nkpolar): + pn = np.where((kp < akplim[ik + 1]) & (kp >= akplim[ik])) + polar_index[pn[0], pn[1], pn[2]] = ik + nbin[ik] = nbin[ik] + len(pn[0]) + + return akp, nbin, polar_index, akplim + + +def polar_isotropic(nkpolar, nkx, nky, nkz, polar_index, nbin, fft_matrix, kx, + ky, kz): + """Average a spectrum over polar (k-perpendicular) shells. + + Accumulates the values of ``fft_matrix`` into the polar bins defined by + ``polar_index`` (as produced by :func:`init_polar`) and divides by the + number of cells per bin to obtain the isotropic (shell-averaged) + spectrum. Works for 2D grids (set ``nkz`` and ``kz`` to ``0``) and 3D + grids. + + Args: + nkpolar: Number of polar (radial ``k_perp``) bins. + nkx: Number of grid points along the ``kx`` axis. + nky: Number of grid points along the ``ky`` axis. + nkz: Number of grid points along the ``kz`` axis; use ``0`` for 2D data. + polar_index: Integer array mapping each Cartesian wavenumber cell to + its polar bin, as returned by :func:`init_polar`. + nbin: Number of Cartesian cells in each polar bin, used as the + averaging denominator. + fft_matrix: Spectral quantity (e.g. spectral power) defined on the + Cartesian wavenumber grid to be averaged over shells. + kx: 1D array of ``kx`` wavenumbers (accepted for interface consistency). + ky: 1D array of ``ky`` wavenumbers (accepted for interface consistency). + kz: 1D array of ``kz`` wavenumbers (accepted for interface consistency). + + Returns: + The shell-averaged (isotropic) spectrum, one value per polar bin + (shape ``(nkpolar,)``). + """ + # if 2D, then nkz = kz = 0 + + fft_isok = np.zeros(nkpolar) + if nkz == 0: + for i in range(nkx): + for j in range(nky): + fft_isok[polar_index[i, + j]] = fft_isok[polar_index[i, j]] + fft_matrix[i, + j] + else: + for i in range(nkx): + for j in range(nky): + for k in range(nkz): + fft_isok[polar_index[ + i, j, k]] = fft_isok[polar_index[i, j, k]] + fft_matrix[i, j, k] + + fft_isok = fft_isok / nbin[:] + return fft_isok diff --git a/src/postgkyl/numerics/filters.py b/src/postgkyl/numerics/filters.py new file mode 100644 index 00000000..0fea7933 --- /dev/null +++ b/src/postgkyl/numerics/filters.py @@ -0,0 +1,73 @@ +"""Low-pass filtering: FFT brick-wall and Butterworth. + +The legacy ``tools/filters.py`` fell back to an interactive matplotlib +click-to-pick cutoff frequency when ``cutoff`` was omitted. That picker is +an effect at the edge (it pops up a figure and blocks on a GUI event) and +does not belong in a pure-array leaf module; it has not been ported here. +If anyone still wants that convenience, it belongs in ``render``/``cli``, +built on top of :func:`fft_filtering`. Consequently ``cutoff`` is a +required argument here rather than optional. +""" + +from __future__ import annotations + +import numpy as np +from scipy.signal import butter, lfilter + + +def fft_filtering(data: np.ndarray, + dt: float = 1.0, + *, + cutoff: float) -> np.ndarray: + """Low-pass filter ``data`` by zeroing FFT bins above ``cutoff``. + + Args: + data: 1-D signal. + dt: Sample spacing. + cutoff: High-frequency cutoff; bins with ``|freq| > cutoff`` are zeroed. + + Returns: + The (complex) inverse FFT of the filtered spectrum. + """ + N = len(data) + freq = np.fft.fftfreq(N, dt) + FT = np.fft.fft(data) + + FT[freq > cutoff] = 0 + FT[freq < -cutoff] = 0 + + return np.fft.ifft(FT) + + +def _butter_lowpass(cutoff: float, fs: float, order: int = 5): + nyq = 0.5 * fs + normal_cutoff = cutoff / nyq + b, a = butter(order, normal_cutoff, btype="low", analog=False) + return b, a + + +def _butter_lowpass_filter(data: np.ndarray, + cutoff: float, + fs: float, + order: int = 5): + b, a = _butter_lowpass(cutoff, fs, order=order) + return lfilter(b, a, data) + + +def butter_filtering(data: np.ndarray, + dt: float = 1.0, + *, + cutoff: float) -> np.ndarray: + """Low-pass filter ``data`` with a 6th-order Butterworth filter. + + Args: + data: 1-D signal. + dt: Sample spacing. + cutoff: High-frequency cutoff. + + Returns: + The filtered signal (same length as ``data``). + """ + order = 6 + fs = 1 / dt # sample rate + return _butter_lowpass_filter(data, cutoff, fs, order) diff --git a/src/postgkyl/numerics/fit.py b/src/postgkyl/numerics/fit.py new file mode 100644 index 00000000..7d06a28b --- /dev/null +++ b/src/postgkyl/numerics/fit.py @@ -0,0 +1,458 @@ +"""Curve fitting: built-in model functions (including ``exp2``, the +growth-rate model), an RPN custom-model parser, ``scipy.optimize.curve_fit`` +wrappers, and the leading-window search used for growth-rate-style fits.""" + +from __future__ import annotations + +from typing import Callable + +import numpy as np +import scipy.optimize as opt + + +def linear(x: np.ndarray, a: float, b: float) -> np.ndarray: + """``f(x) = a*x + b`` + + a: slope (change in f per unit x). + b: intercept, f(0). + """ + return a * x + b + + +def quadratic(x: np.ndarray, a: float, b: float, c: float) -> np.ndarray: + """``f(x) = a*x**2 + b*x + c`` + + a: quadratic coefficient; the second derivative is 2*a. + b: linear coefficient, the slope at x = 0. + c: intercept, f(0). + """ + return a * x**2 + b * x + c + + +def plane(XY: np.ndarray, a: float, b: float, c: float) -> np.ndarray: + """``f(x, y) = a*x + b*y + c`` + + a: slope along x at fixed y. + b: slope along y at fixed x. + c: intercept, f(0, 0). + """ + x, y = XY + return a * x + b * y + c + + +def quadratic2d(XY: np.ndarray, a: float, b: float, c: float, d: float, + e: float, f: float) -> np.ndarray: + """``fitted(x, y) = a*x**2 + b*y**2 + c*x*y + d*x + e*y + f`` + + a: x-squared coefficient. + b: y-squared coefficient. + c: cross-term coefficient multiplying x*y. + d: linear x coefficient. + e: linear y coefficient. + f: intercept, fitted(0, 0). + """ + x, y = XY + return a * x**2 + b * y**2 + c * x * y + d * x + e * y + f + + +def exp_plateau(x: np.ndarray, A: float, b: float, C: float) -> np.ndarray: + """``f(x) = A*exp(b*x) + C`` + + A: initial offset from C; f(0) = A + C. + b: exponential rate (inverse x units); b < 0 means decay toward C. + C: plateau approached as b*x -> -infinity. + """ + return A * np.exp(b * x) + C + + +def gaussian(x: np.ndarray, A: float, mu: float, sigma: float) -> np.ndarray: + """``f(x) = A * exp(-0.5 * ((x - mu) / sigma)**2)`` + + A: amplitude at the center, f(mu). + mu: center position. + sigma: width parameter; abs(sigma) is the standard deviation in x units. + """ + return A * np.exp(-0.5 * ((x - mu) / sigma)**2) + + +def power(x: np.ndarray, a: float, n: float, b: float) -> np.ndarray: + """``f(x) = a * x**n + b`` + + a: amplitude multiplying the power law. + n: power-law exponent. + b: additive offset. + """ + return a * x**n + b + + +def sinusoid(x: np.ndarray, A: float, omega: float, phi: float, + C: float) -> np.ndarray: + """``f(x) = A * sin(omega * x + phi) + C`` + + A: signed oscillation amplitude. + omega: angular frequency (radians per unit x). + phi: phase at x = 0 (radians). + C: mean level of the oscillation. + """ + return A * np.sin(omega * x + phi) + C + + +def tanh_transition(x: np.ndarray, A: float, x0: float, w: float, + C: float) -> np.ndarray: + """``f(x) = A * tanh((x - x0) / w) + C`` + + A: signed half-difference between the two asymptotic levels C - A and C + A. + x0: transition midpoint, where f(x0) = C. + w: transition scale in x units; the slope at x0 is A/w. + C: midpoint level. + """ + return A * np.tanh((x - x0) / w) + C + + +def exp2(x: np.ndarray, a: float, b: float) -> np.ndarray: + """``f(x) = a * exp(2*b*x)`` + + a: initial value, f(0). + b: amplitude growth rate (inverse x units); the fitted curve's rate is 2*b. + Energy (a squared quantity) is typically used for growth-rate studies, + hence the factor of 2 in the exponent. + """ + return a * np.exp(2 * b * x) + + +RPN_OPERATORS: frozenset = frozenset({'+', '-', '*', '/', '**', '^'}) + +RPN_FUNCTIONS: dict[str, Callable] = { + 'exp': np.exp, + 'log': np.log, + 'ln': np.log, + 'log10': np.log10, + 'sin': np.sin, + 'cos': np.cos, + 'tan': np.tan, + 'sqrt': np.sqrt, + 'abs': np.abs, + 'tanh': np.tanh, +} + +_SPATIAL_VARS: frozenset = frozenset({'x', 'y', 'z'}) + + +def rpn_param_names(expression: str) -> list[str]: + """Return the free parameter names from an RPN expression, in order of + first appearance.""" + names = [] + for tok in expression.split(): + if tok in _SPATIAL_VARS or tok in RPN_OPERATORS or tok in RPN_FUNCTIONS: + continue + try: + float(tok) + except ValueError: + if tok not in names: + names.append(tok) + return names + + +def rpn_ndim(expression: str) -> int: + """Return 1 or 2 depending on whether ``y`` appears as a spatial variable.""" + return 2 if 'y' in expression.split() else 1 + + +def _rpn_make_func(expression: str) -> Callable: + """Build a ``curve_fit``-compatible callable from an RPN expression string.""" + tokens = expression.split() + param_names = rpn_param_names(expression) + ndim = rpn_ndim(expression) + + def _func(xdata, *param_values): + ns: dict = dict(zip(param_names, param_values)) + if ndim == 1: + ns['x'] = np.asarray(xdata, dtype=float) + else: + ns['x'] = np.asarray(xdata[0], dtype=float) + ns['y'] = np.asarray(xdata[1], dtype=float) + + stack = [] + for tok in tokens: + if tok in RPN_OPERATORS: + b, a = stack.pop(), stack.pop() + if tok == '+': + stack.append(a + b) + elif tok == '-': + stack.append(a - b) + elif tok == '*': + stack.append(a * b) + elif tok == '/': + stack.append(a / b) + else: + stack.append(a**b) # ** or ^ + elif tok in RPN_FUNCTIONS: + stack.append(RPN_FUNCTIONS[tok](stack.pop())) + elif tok in ns: + stack.append(ns[tok]) + else: + stack.append(float(tok)) + + result = stack[0] + ref = ns.get('x', ns.get('y')) + if np.ndim(result) == 0 and ref is not None: + result = np.full_like(ref, float(result)) + return np.asarray(result, dtype=float) + + return _func + + +FIT_FUNCTIONS: dict[str, Callable] = { + "linear": linear, + "quadratic": quadratic, + "plane": plane, + "quadratic2d": quadratic2d, + "exp_plateau": exp_plateau, + "gaussian": gaussian, + "power": power, + "sinusoid": sinusoid, + "tanh_transition": tanh_transition, + "exp2": exp2, +} + +# Number of spatial dimensions each fit type operates on +FIT_NDIM: dict[str, int] = { + "linear": 1, + "quadratic": 1, + "plane": 2, + "quadratic2d": 2, + "exp_plateau": 1, + "gaussian": 1, + "power": 1, + "sinusoid": 1, + "tanh_transition": 1, + "exp2": 1, +} + + +def fit_evaluate(xdata: np.ndarray, fit_type: str, + params: np.ndarray) -> np.ndarray: + """Evaluate a fitted model at ``xdata`` given the optimized parameters.""" + if fit_type in FIT_FUNCTIONS: + return FIT_FUNCTIONS[fit_type](xdata, *params) + return _rpn_make_func(fit_type)(xdata, *params) + + +def fit(xdata: np.ndarray, + ydata: np.ndarray, + fit_type: str = "linear", + p0: list | None = None) -> tuple[np.ndarray, np.ndarray, float]: + """Fit data using ``scipy.optimize.curve_fit`` with the specified model. + + Args: + xdata: For 1D fits, shape ``(N,)``. For 2D fits, shape ``(2, N)`` where + rows are the two independent variables flattened. + ydata: Dependent variable, shape ``(N,)``. + fit_type: A key in :data:`FIT_FUNCTIONS`, or an RPN expression string + (e.g. ``"a x * b +"``). + p0: Initial guess for the fit parameters; defaults to all ones. + + Returns: + ``(params, cov, R2)``. + + Raises: + ValueError: If ``fit_type`` is neither a known model name nor a + recognizable RPN expression. + """ + if fit_type in FIT_FUNCTIONS: + func = FIT_FUNCTIONS[fit_type] + n_params = func.__code__.co_argcount - 1 + else: + toks = set(fit_type.split()) + if not (toks & (RPN_OPERATORS | set(RPN_FUNCTIONS))): + raise ValueError( + f"fit_type '{fit_type}' not recognized. Choose from: {list(FIT_FUNCTIONS)}" + ) + func = _rpn_make_func(fit_type) + n_params = len(rpn_param_names(fit_type)) + + if p0 is None: + p0 = np.ones(n_params) + + params, cov = opt.curve_fit(func, xdata, ydata, p0=p0) + + residual = ydata - func(xdata, *params) + ss_res = np.sum(residual**2) + ss_tot = np.sum((ydata - np.mean(ydata))**2) + R2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else 1.0 + + return params, cov, R2 + + +def auto_guess(fit_type: str, xdata: np.ndarray, + ydata: np.ndarray) -> list | None: + """Return data-driven initial parameter guesses for known fit types. + + Produces a sensible ``p0`` for :func:`fit` by inspecting the data (e.g. a + least-squares seed for linear/polynomial models, peak location and FWHM + for a gaussian, the dominant FFT frequency for a sinusoid). Returns + ``None`` for RPN expressions or when the data has no finite values, in + which case :func:`fit` falls back to its default (ones). + + Args: + fit_type: A built-in model name (an RPN expression yields ``None``). + xdata: Independent variable: shape ``(N,)`` for 1D models, ``(2, N)`` + for 2D. + ydata: Dependent variable, shape ``(N,)``. + + Returns: + A list of initial parameter guesses, or ``None`` when no heuristic + applies. + """ + y = np.asarray(ydata, dtype=float) + finite = np.isfinite(y) + if not np.any(finite): + return None + y_fin = y[finite] + y_min, y_max = y_fin.min(), y_fin.max() + y_mean = y_fin.mean() + y_range = y_max - y_min + + if fit_type == "linear": + x = np.asarray(xdata) + dx = x.max() - x.min() + a = y_range / dx if dx != 0 else 1.0 + b = y_mean - a * x.mean() + return [a, b] + + if fit_type == "quadratic": + x = np.asarray(xdata) + try: + return list(np.polyfit(x, y, 2)) + except Exception: + return [0.0, 1.0, y_mean] + + if fit_type == "plane": + x, yc = xdata[0], xdata[1] + A = np.column_stack([x, yc, np.ones_like(x)]) + result, *_ = np.linalg.lstsq(A, y, rcond=None) + return list(result) + + if fit_type == "quadratic2d": + x, yc = xdata[0], xdata[1] + A = np.column_stack([x**2, yc**2, x * yc, x, yc, np.ones_like(x)]) + result, *_ = np.linalg.lstsq(A, y, rcond=None) + return list(result) + + if fit_type == "exp_plateau": + x = np.asarray(xdata) + n_tail = max(1, len(x) // 10) + C = float(y[np.argsort(x)[-n_tail:]].mean()) + A = float(y_max - C) or 1.0 + x_span = x.max() - x.min() + b = -1.0 / x_span if x_span > 0 else -1.0 + return [A, b, C] + + if fit_type == "gaussian": + x = np.asarray(xdata) + A = float(y_max) + mu = float(x[np.argmax(y)]) + above = x[y >= A / 2] if A != 0 else x + if len(above) >= 2: + sigma = float((above[-1] - above[0]) / (2 * np.sqrt(2 * np.log(2)))) + else: + sigma = float((x.max() - x.min()) / 4) + return [A, mu, max(abs(sigma), 1e-10)] + + if fit_type == "power": + b_off = float(y_min) + a = float(y_max - b_off) or 1.0 + return [a, 1.0, b_off] + + if fit_type == "sinusoid": + x = np.asarray(xdata) + A = float(y_range / 2) or 1.0 + C = float((y_max + y_min) / 2) + sort_idx = np.argsort(x) + x_s, y_s = x[sort_idx], y[sort_idx] + if len(x_s) > 1: + dx = np.mean(np.diff(x_s)) + freqs = np.fft.rfftfreq(len(y_s), d=dx) + fft_amp = np.abs(np.fft.rfft(y_s - C)) + i_peak = np.argmax(fft_amp[1:]) + 1 if len(fft_amp) > 1 else 1 + omega = float(2 * np.pi * freqs[i_peak]) + else: + omega = 1.0 + return [A, omega, 0.0, C] + + if fit_type == "tanh_transition": + x = np.asarray(xdata) + A = float(y_range / 2) or 1.0 + C = float((y_max + y_min) / 2) + x0 = float(x[np.argmax(np.abs(np.gradient(y)))]) + w = float((x.max() - x.min()) / 4) or 1.0 + return [A, x0, w, C] + + if fit_type == "exp2": + # log(y) = log(a) + 2*b*x is linear -- a log-linear regression gives a + # scale-invariant guess without needing to normalize x for curve_fit. + x = np.asarray(xdata, dtype=float) + y_pos = np.clip(y, 1e-300, None) + slope, intercept = np.polyfit(x, np.log(y_pos), 1) + return [float(np.exp(intercept)), float(slope / 2)] + + return None + + +def fit_best_window( + xdata: np.ndarray, + ydata: np.ndarray, + fit_type: str = "exp2", + min_n: int | None = None, + p0: list | None = None) -> tuple[np.ndarray, np.ndarray, float, int]: + """Fit ``fit_type`` to the best-scoring leading window of a 1D series. + + Scans windows ``xdata[:n]`` for ``n`` from ``min_n`` up to ``len(xdata)``, + keeping the window with the best coefficient of determination (R^2). Each + window is warm-started from the previous window's fitted parameters (or + ``p0``/:func:`auto_guess` for the first), so this generalizes a single + full-domain :func:`fit` call to the common case of a time series whose + early or late region should be excluded (e.g. growth-rate fits, which are + only valid while the signal grows/decays continuously). + + Args: + xdata: 1D independent variable (e.g. time). + ydata: dependent variable, shape matching ``xdata``. + fit_type: passed to :func:`fit`. + min_n: minimum number of points in the fitted window. Defaults to + ``len(xdata) // 10``. + p0: initial guess for the first window; ``None`` uses :func:`auto_guess`. + + Returns: + ``(params, cov, R2, N)`` for the best-scoring window. + + Raises: + RuntimeError: if ``curve_fit`` fails to converge for every window in + the scan range. + """ + xdata = np.asarray(xdata, dtype=float) + ydata = np.asarray(ydata, dtype=float) + if min_n is None: + min_n = max(2, len(xdata) // 10) + + best_R2 = -np.inf + best = None + guess = p0 + for n in range(min_n, len(xdata) + 1): + xn, yn = xdata[:n], ydata[:n] + try: + params, cov, R2 = fit( + xn, + yn, + fit_type, + p0=guess if guess is not None else auto_guess(fit_type, xn, yn)) + except RuntimeError: + continue + guess = list(params) + if R2 > best_R2: + best_R2, best = R2, (params, cov, R2, n) + if best is None: + raise RuntimeError( + "fit_best_window: curve_fit failed to converge for every window in " + f"[{min_n:d}, {len(xdata):d}]") + return best diff --git a/src/postgkyl/numerics/grid_centering.py b/src/postgkyl/numerics/grid_centering.py new file mode 100644 index 00000000..1e968a1f --- /dev/null +++ b/src/postgkyl/numerics/grid_centering.py @@ -0,0 +1,52 @@ +"""Convert a nodal (edge) grid to its cell-centered equivalent.""" + +from __future__ import annotations + +import numpy as np + + +def nodal_to_cell_centered_grid(grid: list[np.ndarray], + cells: np.ndarray, + meshgrid: bool = False) -> list[np.ndarray]: + """Return the cell-centered grid corresponding to a nodal (edge) grid. + + Args: + grid: List of NumPy arrays giving the nodal grid coordinates. + cells: Number of cells in each dimension. + meshgrid: If ``True`` and the coordinates are 1-D, return an + ij-indexed meshgrid instead of the plain 1-D per-axis arrays. + + Returns: + List of NumPy arrays giving the cell-centered grid coordinates. + + Raises: + ValueError: If ``grid`` and ``cells`` disagree on the number of + dimensions, or an axis is neither nodal nor already cell-centered. + """ + num_dims = len(grid) + grid_out = [] + if num_dims != len(cells): + raise ValueError("Number dimensions for 'grid' and 'values' doesn't match") + for d in range(num_dims): + if len(grid[d].shape) == 1: + if grid[d].shape[0] == cells[d]: + grid_out.append(grid[d]) + elif grid[d].shape[0] == cells[d] + 1: + grid_out.append(0.5 * (grid[d][:-1] + grid[d][1:])) + else: + raise ValueError("Something is terribly wrong...") + else: + if grid[d].shape[d] == cells[d]: + grid_out.append(grid[d]) + elif grid[d].shape[d] == cells[d] + 1: + if num_dims == 1: + grid_out.append(0.5 * (grid[d][:-1] + grid[d][1:])) + else: + grid_out.append(0.5 * (grid[d][:-1, :-1] + grid[d][1:, 1:])) + else: + raise ValueError("Something is terribly wrong...") + + if meshgrid and num_dims > 1 and all(axis.ndim == 1 for axis in grid_out): + return list(np.meshgrid(*grid_out, indexing="ij")) + + return grid_out diff --git a/src/postgkyl/numerics/idx_parser.py b/src/postgkyl/numerics/idx_parser.py new file mode 100644 index 00000000..2df89d96 --- /dev/null +++ b/src/postgkyl/numerics/idx_parser.py @@ -0,0 +1,71 @@ +"""Parse index / value / slice selectors into NumPy indices (pure).""" + +from __future__ import annotations + +import numpy as np + + +def _find_nearest_index(array, value): + if array is None: + raise TypeError( + "Float selector given but no coordinate array to match against.") + idx = np.searchsorted(array, value) + if idx == len(array): + return int(idx - 2) + elif idx > 0: + return int(idx - 1) + else: + return int(idx) + + +def _find_cell_index(array, value): + if array is None: + raise TypeError( + "Float selector given but no coordinate array to match against.") + return int(np.searchsorted(array, value)) + + +def _string_to_index(value: str, array: np.ndarray, nodal: bool = False) -> int: + if not isinstance(value, str): + raise TypeError("Value is not a string") + if value.lstrip("-").isdigit(): + return int(value) + return _find_cell_index(array, + float(value)) if nodal else _find_nearest_index( + array, float(value)) + + +def idx_parser(value: int | float | str, + array: np.ndarray | None = None, + nodal: bool = False) -> int | slice | tuple: + """Turn an int/float/str selector into an int index, ``slice``, or tuple. + + - int -> used as-is + - float -> nearest (or containing, if ``nodal``) cell index + - ``"a,b,c"`` -> tuple of indices + - ``"a:b"`` -> ``slice`` + - ``"a"`` -> single index + """ + if isinstance(value, int): + return value + if isinstance(value, float): + return _find_cell_index(array, value) if nodal else _find_nearest_index( + array, value) + if isinstance(value, str): + if len(value.split(",")) > 1: + return tuple(_string_to_index(i, array, nodal) for i in value.split(",")) + if len(value.split(":")) == 2: + lo, hi = value.split(":") + if lo == "": + lo = "0" + if hi == "": + hi = str(len(array)) + try: + if int(hi) < 0: + hi = str(len(array) + int(hi) + 1) + except ValueError: + pass + return slice(_string_to_index(lo, array, nodal), + _string_to_index(hi, array, nodal)) + return _string_to_index(value, array, nodal) + raise TypeError(f"Unsupported selector type: {type(value)!r}") diff --git a/src/postgkyl/numerics/mag_sq.py b/src/postgkyl/numerics/mag_sq.py new file mode 100644 index 00000000..e921d489 --- /dev/null +++ b/src/postgkyl/numerics/mag_sq.py @@ -0,0 +1,27 @@ +"""Magnitude-squared of a (sub-range of) vector-valued field.""" + +from __future__ import annotations + +import numpy as np + + +def mag_sq(grid: list[np.ndarray], + values: np.ndarray, + coords: str = "0:3") -> tuple[list[np.ndarray], np.ndarray]: + """Compute the magnitude squared of a vector field. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension. + values: Data array whose last axis is components. + coords: ``"start:end"`` slice of the component axis to sum the squares + of. Defaults to the first three components (the common + three-component-vector case). + + Returns: + ``(grid, values)`` where ``values`` has the summed components replaced + by a single trailing component (magnitude squared). + """ + lo, hi = coords.split(":") + comps = values[..., slice(int(lo), int(hi))] + out = np.sum(comps * comps, axis=-1)[..., np.newaxis] + return list(grid), out diff --git a/src/postgkyl/numerics/natural_sort.py b/src/postgkyl/numerics/natural_sort.py new file mode 100644 index 00000000..cba5b8c8 --- /dev/null +++ b/src/postgkyl/numerics/natural_sort.py @@ -0,0 +1,14 @@ +"""Natural/numeric sort key for strings (pure).""" + +from __future__ import annotations + +import re + +_CHUNK_RE = re.compile(r"(\d+)") + + +def natural_sort_key(s: str) -> tuple: + """Split ``s`` into alternating text/int chunks so embedded digit runs + compare numerically instead of character-by-character (``'field_2'`` sorts + before ``'field_10'``, unlike a plain lexicographic string sort).""" + return tuple(int(c) if c.isdigit() else c for c in _CHUNK_RE.split(s)) diff --git a/src/postgkyl/numerics/rel_change.py b/src/postgkyl/numerics/rel_change.py new file mode 100644 index 00000000..b3eaca33 --- /dev/null +++ b/src/postgkyl/numerics/rel_change.py @@ -0,0 +1,30 @@ +"""Relative change of one dataset's values against a reference.""" + +from __future__ import annotations + +import numpy as np + + +def rel_change(grid: list[np.ndarray], + values0: np.ndarray, + values: np.ndarray, + comp: int | None = None) -> tuple[list[np.ndarray], np.ndarray]: + """Compute ``(values - values0) / values0``, component-wise. + + Args: + grid: Nodal coordinate arrays, one per spatial dimension (returned + unchanged; the two datasets are assumed to share a grid). + values0: Reference ("before") data array. + values: Data array to compare against the reference. + comp: If given, every component is normalized by this single reference + component instead of its own (e.g. divide every energy component by + the total energy component). + + Returns: + ``(grid, out)`` with ``out`` the same shape as ``values``. + """ + out = np.zeros(values.shape) + for i in range(out.shape[-1]): + denom = values0[..., int(comp)] if comp is not None else values0[..., i] + out[..., i] = (values[..., i] - values0[..., i]) / denom + return list(grid), out diff --git a/src/postgkyl/numerics/rotation_matrix.py b/src/postgkyl/numerics/rotation_matrix.py new file mode 100644 index 00000000..35a4e40c --- /dev/null +++ b/src/postgkyl/numerics/rotation_matrix.py @@ -0,0 +1,34 @@ +"""Rotation matrix aligning the x-axis with a given vector.""" + +from __future__ import annotations + +import numpy as np + + +def rotation_matrix(vector: np.ndarray) -> np.ndarray: + """Calculate a 3x3 rotation matrix whose first row is ``vector``'s direction. + + Args: + vector: A 3-component vector (nonzero in every component). + + Returns: + 3x3 rotation matrix (NumPy array). + """ + rot = np.zeros((3, 3)) + norm = np.abs(vector) + k = vector / norm # direction unit vector + + # normalization + norm2 = np.sqrt(k[1] * k[1] + k[2] * k[2]) + norm3 = np.sqrt((k[1] * k[1] + k[2] * k[2])**2 + k[0] * k[0] * k[1] * k[1] + + k[0] * k[0] * k[2] * k[2]) + + rot[0, :] = k + rot[1, 0] = 0 + rot[1, 1] = -k[2] / norm2 + rot[1, 2] = k[1] / norm2 + rot[2, 0] = (k[1] * k[1] + k[2] * k[2]) / norm3 + rot[2, 1] = -k[0] * k[1] / norm3 + rot[2, 2] = -k[0] * k[2] / norm3 + + return rot diff --git a/src/postgkyl/operations/__init__.py b/src/postgkyl/operations/__init__.py new file mode 100644 index 00000000..e07afcf4 --- /dev/null +++ b/src/postgkyl/operations/__init__.py @@ -0,0 +1,181 @@ +"""The data-transformation library -- one function per operation. + +Every verb takes a dataset first and returns a dataset (via ``_result``), so the +fluent ``GData`` methods, the operators, and any CLI all delegate here and can +never drift apart. Verbs are typed on ``GDataState`` but return the caller's +concrete (sub)class because ``_result`` rebuilds ``type(self)``. + +``interpolate`` is the one-way modal -> NumPy bridge; ``arithmetic`` dispatches +on the container backend (Gkeyll kernels for modal data, NumPy for field data); +``integrate`` performs full or partial integration inside Gkeyll on modal +data (full is terminal; partial stays native and lower-dimensional); +``average`` reduces modal data over a dimension subset via +``gkyl_array_average``, producing a new lower-dimensional modal dataset; +``map`` delegates to the grid-mapping engine in ``dg.map``. Flat modules are +domain-independent core verbs; domain subpackages such as ``gyrokinetics`` +hold transformations that require domain geometry without interpreting field +components as new physical conclusions. Equation-specific physics (the former +``moments``/``agyro``/``current``/``energetics``/``rotate``/ +``transform_frame``/``laguerre`` verbs, folded with the array math they +delegated to) lives one layer up, in ``diagnostics``. + +The terminal renderers (``plot``, ``animate``, ``plotly``, ``plotly_animate``, +and ``pyvista``) are exceptions: +this namespace re-exports their exact canonical callables from +:mod:`postgkyl.render` without wrapping them. +""" + +from . import arithmetic, gyrokinetics +from .interpolate import interpolate +from .local_poly import local_poly +from .select import select +from .info import info +from .print import print +from .integrate import integrate +from .average import average +from .eval_at_coord_proj import eval_at_coord_proj +from postgkyl.render import animate, plot, plotly, plotly_animate, pyvista +from .represent import apply, represent + +from .fft import fft +from .magsq import magsq +from .relchange import relchange +from .mask import mask +from .collect import collect +from .sort import sort +from .grid import grid +from .val2coord import val2coord +from .extract_input import extract_input +from .fit import fit +from .growth import growth +from .differentiate import differentiate +from .evaluate import available_operators as available_evaluate_operators, evaluate +from .map import map + +# Command metadata is attached at the layer that owns each operation. This +# block is deliberately declarative: discovery still walks the public API and +# there is no registration side effect or CLI import here. +from typing import Annotated, Literal + +from postgkyl.cli_spec import ( + CliArgument, + CliType, + CommandSpec, + DatasetRef, + Execution, + ResultPolicy, + Section, + command, + hidden, +) +from postgkyl.gdatastate.gdatastate import GDataState + + +def _resolve_receiver_annotations(*functions) -> None: + """Make the modules' public forward references runtime-resolvable.""" + for function in functions: + function.__globals__.setdefault("GDataState", GDataState) + function.__globals__.setdefault("_GDataState", GDataState) + + +_MAP = CommandSpec(Section.VERBS, Execution.MAP_REPLACE) +_APPEND = CommandSpec(Section.VERBS, Execution.MAP_APPEND, consumes_inputs=True) +_COMBINE = CommandSpec(Section.VERBS, Execution.COMBINE, consumes_inputs=True) +_TERM_EACH = CommandSpec(Section.UTILITY, + Execution.TERMINAL_EACH, + result=ResultPolicy.VALUE) +_TERM_ALL = CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.VALUE) + +for _function in ( + interpolate, + local_poly, + select, + integrate, + average, + eval_at_coord_proj, + fft, + magsq, + relchange, + mask, + collect, + sort, + grid, + val2coord, + extract_input, + fit, + differentiate, + evaluate, + map, + represent, + growth, +): + _resolve_receiver_annotations(_function) + +select.__annotations__.update(comp=str | None, + z0=str | None, + z1=str | None, + z2=str | None, + z3=str | None, + z4=str | None, + z5=str | None) +integrate.__annotations__["op"] = Literal["none", "abs", "sq"] +integrate.__annotations__["axis"] = Annotated[int | tuple | str | None, + CliType(str | None), + CliArgument()] +evaluate.__annotations__["chain"] = Annotated[str, CliArgument()] +average.__annotations__["dims"] = list[int] +average.__annotations__["weight"] = Annotated[GDataState | None, DatasetRef()] +eval_at_coord_proj.__annotations__.update(eval_dirs=list[int], + eval_coords=list[float]) +relchange.__annotations__.update(data0=Annotated[GDataState, + DatasetRef()], + data=Annotated[GDataState, + DatasetRef()], + comp=str | None) +mask.__annotations__["mask_data"] = Annotated[GDataState | None, DatasetRef()] +fit.__annotations__["guess"] = str | None +map.__annotations__["data"] = GDataState +map.__annotations__["mapping"] = str +represent.__annotations__["to"] = Literal["modal", "nodal", "quad"] + +for _function in (interpolate, local_poly, select, average, eval_at_coord_proj, + fft, magsq, grid, differentiate, map): + command(_MAP)(_function) +command(_APPEND)(val2coord) +command(_COMBINE)(relchange) +command(_COMBINE)(collect) +command(CommandSpec(Section.VERBS, Execution.COMBINE, + consumes_inputs=True))(sort) +command(_COMBINE)(evaluate) +command(_MAP)(mask) +command(CommandSpec(Section.VERBS, Execution.MAP_APPEND))(fit) +command(CommandSpec(Section.VERBS, Execution.MAP_APPEND))(growth) +command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT))(info) +command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT))(print) +command( + CommandSpec(Section.VERBS, + Execution.MAP_OR_TERMINAL_EACH, + result=ResultPolicy.VALUE))(integrate) +command(_TERM_EACH)(extract_input) +command(_MAP)(represent) + +hidden("requires a Python callable and cannot be lowered losslessly")(apply) +hidden("registry provider used by evaluate help and validation")( + available_evaluate_operators) + +__all__ = [ + "interpolate", "local_poly", "select", "info", "print", "integrate", + "average", "eval_at_coord_proj", "plot", "animate", "plotly", + "plotly_animate", "pyvista", "arithmetic", "represent", "apply", "fft", + "magsq", "relchange", "mask", "collect", "sort", "grid", "val2coord", + "extract_input", "fit", "differentiate", "evaluate", + "available_evaluate_operators", "map", "growth", "gyrokinetics" +] diff --git a/src/postgkyl/operations/_curvilinear.py b/src/postgkyl/operations/_curvilinear.py new file mode 100644 index 00000000..859561d1 --- /dev/null +++ b/src/postgkyl/operations/_curvilinear.py @@ -0,0 +1,35 @@ +"""Shared "which dimensions form a curvilinear block" lookup. + +``differentiate`` and ``integrate`` (its per-axis mode) both need +``ctx["mapped_axes"]`` grouped into contiguous, genuinely curvilinear +(``ndim > 1`` -- a joint ``space="conf"`` ``.map()``) blocks before they can +apply the chain-rule/Jacobian math in ``numerics.curvilinear``, mirroring +the sibling-grouping ``select``'s curvilinear guard already does inline. +Kept separate from that inline grouping (rather than factored together) +because ``select`` also needs its *separable* (``space="vel"``) siblings +grouped for its own purposes, where the callers here index a block by its +own local axis order and so need it filtered to curvilinear blocks and +sorted. +""" + +from __future__ import annotations + + +def curvilinear_blocks(grid: list, mapped_axes: dict) -> dict: + """``{offset: sorted [absolute dims]}`` for every genuinely curvilinear + (multi-dimensional grid array) block recorded in ``mapped_axes``.""" + blocks: dict = {} + for d, off in mapped_axes.items(): + if grid[d].ndim > 1: + blocks.setdefault(off, []).append(d) + for dims in blocks.values(): + dims.sort() + return blocks + + +def block_for_axis(blocks: dict, axis: int): + """The ``(offset, dims)`` of the block containing ``axis``, or ``None``.""" + for off, dims in blocks.items(): + if axis in dims: + return off, dims + return None diff --git a/src/postgkyl/operations/arithmetic.py b/src/postgkyl/operations/arithmetic.py new file mode 100644 index 00000000..12414623 --- /dev/null +++ b/src/postgkyl/operations/arithmetic.py @@ -0,0 +1,251 @@ +"""Arithmetic / NumPy-ufunc backend for the fluent operators. + +Defined here (in ``operations``) -- not on the container -- so the computing operators +follow the same one-way layering as every other verb (HIERARCHY_3.md). + +Dispatch is on the container's ``backend`` (the two-domain lifecycle of +REFACTOR_GKEYLL_FFI.md): + +- **gkyl-backed (modal) operands** run inside Gkeyll: ``*``/``/`` are the weak + kernels (``gkyl_dg_mul_op``/``div_op``), ``+``/``-`` are coefficient linear + combinations (``gkyl_array_set``/``accumulate``), scalar multiply is + ``gkyl_array_scale``, scalar add shifts the mean coefficient, positive + integer powers are repeated weak multiplies, and any other power (0, + negative, or fractional) is ``gkyl_proj_powsqrt_on_basis`` (a + quadrature projection of ``pow(sqrt(f), 2*exponent)``). Results stay + modal (gkyl-backed). + Two modal operands of *different* dimensionality (e.g. a conf-space density + times a phase-space distribution) automatically route ``*`` through + ``gkyl_dg_mul_conf_phase_op_range`` instead -- whichever operand has fewer + dimensions is the conf side, independent of call order. +- **numpy-backed operands** take the unchanged NumPy path. +- **Mixing the domains** in one expression is an error naming the fix. +""" + +from __future__ import annotations + +import operator + +import numpy as np + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl import dg, numerics + + +def _unpack(x): + """(values, grid, dataset|None) for a dataset; (array, None, None) otherwise.""" + if isinstance(x, GDataState): + return x.values, x.grid, x + return np.asarray(x), None, None + + +def binary(op, a, b): + """``a b`` where at least one operand is a dataset; result copies its grid.""" + pa = a if isinstance(a, GDataState) else None + pb = b if isinstance(b, GDataState) else None + if (pa is not None and pa.backend == "gkyl") or (pb is not None + and pb.backend == "gkyl"): + return _modal_binary(op, a, b, pa, pb) + return _numpy_binary(op, a, b, pa, pb) + + +# --------------------------------------------------------------- numpy domain +def _numpy_binary(op, a, b, pa, pb): + va, ga, _ = _unpack(a) + vb, gb, _ = _unpack(b) + primary = pa if pa is not None else pb + primary._require_operable() + if pa is not None and pb is not None: + pb._require_operable() + if not numerics.grids_compatible(ga, gb): + raise ValueError("operands live on different grids") + if va.shape != vb.shape: + raise ValueError(f"incompatible shapes {va.shape} vs {vb.shape}") + return primary._result(primary.grid, op(va, vb)) + + +# --------------------------------------------------------------- modal domain +def _basis_of(data: GDataState): + """(basis_type, ndim, poly_order) from ctx -- the modal ops' dispatch key.""" + basis_type = data.ctx.get("basis_type") + poly_order = data.ctx.get("poly_order") + if basis_type is None or poly_order is None: + raise ValueError("modal operand has no basis_type/poly_order metadata") + return str(basis_type), data.num_dims, int(poly_order) + + +def _modal_binary(op, a, b, pa, pb): + if pa is not None and pb is not None: + return _modal_dataset_pair(op, pa, pb) + primary = pa if pa is not None else pb + other = b if pa is not None else a + if not isinstance(other, (int, float, np.integer, np.floating)): + raise ValueError( + "cannot mix native modal data with arrays; call .interpolate() on the " + "modal operand first (or use scalars / another modal dataset).") + return _modal_scalar(op, primary, float(other), scalar_first=pa is None) + + +def _rep_of(data: GDataState) -> str: + return data.ctx.get("value_form", "modal") + + +def _modal_dataset_pair(op, pa: GDataState, pb: GDataState): + if pb.backend != "gkyl" or pa.backend != "gkyl": + raise ValueError( + "one operand is modal (gkyl-native) and the other is interpolated; " + "call .interpolate() on the modal operand to combine them.") + if pa.num_dims != pb.num_dims: + return _modal_conf_phase_mul(op, pa, pb) + if not numerics.grids_compatible(pa.grid, pb.grid): + raise ValueError("operands live on different grids") + basis = _basis_of(pa) + if _basis_of(pb) != basis: + raise ValueError("operands have different DG bases") + rep = _rep_of(pa) + if rep != _rep_of(pb): + raise ValueError( + f"operands are in different value_forms ({rep} vs {_rep_of(pb)}); " + "convert one explicitly (.to_modal()/.to_nodal()/.to_quad()).") + A, B = pa.native, pb.native + if op is operator.add: # linear: valid in any rep + out = dg.modal.lincomb(1.0, A, 1.0, B) + elif op is operator.sub: + out = dg.modal.lincomb(1.0, A, -1.0, B) + elif rep != "modal": + # Point values (nodal/quad): every pointwise operation is exact -- compute + # with NumPy on the views, wrap back native, stay in-value_form. + out = dg.rep.wrap(op(np.asarray(pa.values), np.asarray(pb.values))) + elif op in (operator.mul, operator.truediv): + out = (dg.modal.weak_mul if op is operator.mul else dg.modal.weak_div)( + *basis, A, B) + else: + raise ValueError( + f"operation {getattr(op, '__name__', op)} is not defined between two " + "modal datasets; .to_nodal()/.to_quad() for pointwise math.") + return pa._result(pa.grid, out) + + +def _modal_conf_phase_mul(op, pa: GDataState, pb: GDataState): + """``conf * phase`` (either order): the operands have different ``num_dims``, + so Gkeyll's per-cell same-basis ``weak_mul`` cannot apply -- this is the + cross-basis ``gkyl_dg_mul_conf_phase_op_range`` path + (``dg.modal.weak_mul_conf_phase``), which multiplies every phase-space cell + by its corresponding lower-dimensional conf-space cell (e.g. a density + times a distribution function). Automatic: whichever operand has fewer + dimensions is the conf side, regardless of call order (``a * b == b * a``). + """ + if op is not operator.mul: + raise ValueError( + f"operands have different dimensionality ({pa.num_dims}D vs " + f"{pb.num_dims}D); only '*' is defined between a lower-dimensional " + "conf-space field and a higher-dimensional phase-space field " + "(Gkeyll has no cross-basis weak divide/add).") + conf, phase = (pa, pb) if pa.num_dims < pb.num_dims else (pb, pa) + for d in (conf, phase): + if _rep_of(d) != "modal": + raise ValueError( + "conf-space x phase-space multiplication is defined for modal DG " + "coefficients only; .to_modal() first.") + if not numerics.grid_is_prefix(conf.grid, phase.grid): + raise ValueError( + "the lower-dimensional operand's grid is not the leading dimensions " + "of the higher-dimensional operand's grid; they are not the same " + "simulation's conf-space and phase-space grids.") + conf_type, conf_ndim, conf_p = _basis_of(conf) + phase_type, phase_ndim, _ = _basis_of(phase) + out = dg.modal.weak_mul_conf_phase(conf_type, conf_ndim, phase_type, + phase_ndim, conf_p, conf.num_cells, + phase.num_cells, conf.native, phase.native) + return phase._result(phase.grid, out) + + +def _modal_scalar(op, data: GDataState, s: float, *, scalar_first: bool): + basis = _basis_of(data) + rep = _rep_of(data) + A = data.native + # Adding/subtracting a *scalar* only shifts the mean (constant) DG + # coefficient -- a constant has no projection onto the higher-order basis + # functions, so gkyl_array_shiftc touches just coefficient 0 (shift_mean, + # dg/modal.py). This is unrelated to array + array (lincomb, above), which + # runs Gkeyll's own accumulate over every coefficient, higher orders + # included. In point-value forms (nodal/quad) there's no separate mean + # coefficient to single out, so a scalar shift moves every component. + shift = (dg.modal.shift_all if rep != "modal" else + lambda a, v: dg.modal.shift_mean(*basis, a, v)) + if op is operator.mul: # linear: valid in any rep + out = dg.modal.scale(A, s) + elif op is operator.truediv and not scalar_first: + out = dg.modal.scale(A, 1.0 / s) # f / s: linear, any rep + elif op is operator.add: + out = shift(A, s) + elif op is operator.sub: + if scalar_first: # s - f + out = shift(dg.modal.scale(A, -1.0), s) + else: # f - s + out = shift(A, -s) + elif rep != "modal": + # Point values: any remaining scalar operation is exact pointwise. + args = (s, np.asarray(data.values)) if scalar_first else (np.asarray( + data.values), s) + out = dg.rep.wrap(op(*args)) + elif op is operator.truediv: # s / f -- weak reciprocal + out = dg.modal.scale(dg.modal.weak_inv(*basis, A), s) + elif op is operator.pow and not scalar_first: + out = dg.modal.power(*basis, + A, + s if not float(s).is_integer() else int(s), + cells=data.ctx.get("cells")) + else: + raise ValueError( + f"operation {getattr(op, '__name__', op)} is not defined for modal " + "data and a scalar; .to_nodal()/.to_quad() for pointwise math.") + return data._result(data.grid, out) + + +# ------------------------------------------------------------------- ufuncs +def apply_ufunc(ufunc, method, *inputs, **kwargs): + """Backend for ``GData.__array_ufunc__``. + + Ufuncs are pointwise, so they are valid wherever the data are point values: + the NumPy field domain, and the nodal/quad value_forms (computed on the + views, wrapped back native, staying in-value_form). Modal coefficients + refuse (via ``_require_operable``): a ufunc has no basis-space meaning. + + Pointwise calls keep the result as a dataset. Reductions return the NumPy + scalar/array produced by the ufunc: after an arbitrary axis reduction the + original spatial grid no longer necessarily describes the result. This + supports NumPy's reduction helpers (``max``, ``min``, ``sum``, ``prod``, + ``all``, and ``any``), which dispatch here as ``ufunc.reduce``. + """ + if method == "reduce": + if len(inputs) != 1 or not isinstance(inputs[0], GDataState): + return NotImplemented + data = inputs[0] + data._require_operable() + return ufunc.reduce(np.asarray(data.values), **kwargs) + if method != "__call__" or "out" in kwargs: + return NotImplemented + primary = next(x for x in inputs if isinstance(x, GDataState)) + primary._require_operable() + rep = (_rep_of(primary) if primary.backend == "gkyl" else None) + raw = [] + for x in inputs: + if isinstance(x, GDataState): + x._require_operable() + if x.backend == "gkyl" and _rep_of(x) != rep or (x.backend != "gkyl" + and rep is not None): + raise ValueError("operands are in different value_forms; convert one " + "explicitly (.to_modal()/.to_nodal()/.to_quad()).") + if x.values.shape != primary.values.shape: + raise ValueError( + f"incompatible shapes {x.values.shape} vs {primary.values.shape}") + raw.append(np.asarray(x.values)) + elif isinstance(x, GDataState._HANDLED_TYPES): + raw.append(x) + else: + return NotImplemented + result = ufunc(*raw, **kwargs) + if rep is not None: + return primary._result(primary.grid, dg.rep.wrap(result)) + return primary._result(primary.grid, result) diff --git a/src/postgkyl/operations/average.py b/src/postgkyl/operations/average.py new file mode 100644 index 00000000..44dd523e --- /dev/null +++ b/src/postgkyl/operations/average.py @@ -0,0 +1,112 @@ +"""The ``average`` verb -- weighted (or plain) average of a native DG field +over a subset of dimensions, via Gkeyll's ``gkyl_array_average``. + +Terminal-adjacent (like ``represent``): unlike ``integrate`` (whose whole-grid +mode returns numbers), this produces a new, lower-dimensional dataset -- still +modal and gkyl-native -- so it composes with ``.to_nodal()``/``.interpolate()``/ +further ``.average()`` calls. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl import dg + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _native_basis(data: "GDataState", what: str): + if data.backend != "gkyl": + raise ValueError( + f"average wraps gkyl_array_average and needs native modal data; " + f"{what} is not available after .interpolate() or without the " + "Gkeyll library.") + if data.ctx.get("value_form", "modal") != "modal": + raise ValueError( + f"average expects the modal value_form, not " + f"'{data.ctx['value_form']}' ({what}); call .to_modal() first.") + basis_type = data.ctx.get("basis_type") + poly_order = data.ctx.get("poly_order") + if basis_type is None or poly_order is None: + raise ValueError(f"{what} has no basis_type/poly_order metadata") + return str(basis_type), int(poly_order) + + +def average(data: "GDataState", + dims, + *, + weight: "GDataState | None" = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """``int f w dx^dims / int w dx^dims`` over the directions in ``dims``. + + Args: + data: gkyl-backed (native modal) dataset in the modal value_form. + dims: iterable of 0-based direction indices to average over (e.g. the + selected ``z0``-``z5`` flags at the CLI layer). + weight: optional gkyl-backed dataset in the modal value_form, same + ``num_dims``/``basis_type``/``poly_order`` as ``data`` and exactly one + field (``gkyl_array_average`` takes no field-index argument) -- the + plain average (dividing by volume) is computed when omitted. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A new dataset over the surviving dimensions -- or a single degenerate + dimension (``grid=[0, 1]``) when every direction is averaged out, + Gkeyll's own convention since there is no true 0-dimensional basis -- + still modal and gkyl-native. + + Raises: + ValueError: ``data`` (or ``weight``) is NumPy-backed or non-modal, is + missing basis metadata, or ``weight``'s grid/basis doesn't match + ``data``'s. + """ + basis_type, poly_order = _native_basis(data, "data") + ndim = data.num_dims + + weight_native = None + if weight is not None: + w_basis_type, w_poly_order = _native_basis(weight, "weight") + if weight.num_dims != ndim: + raise ValueError( + f"weight has {weight.num_dims} dims but the field has {ndim}") + if w_basis_type != basis_type: + raise ValueError( + f"weight basis_type '{w_basis_type}' != field's '{basis_type}'") + if w_poly_order != poly_order: + raise ValueError( + f"weight poly_order {w_poly_order} != field's {poly_order}") + weight_native = weight.native + + grid = { + "ndim": ndim, + "lower": np.asarray(data.ctx["lower"]), + "upper": np.asarray(data.ctx["upper"]), + "cells": np.asarray(data.ctx["cells"]), + } + keep_dirs, cells_avg, out_native = dg.modal.average(grid, + basis_type, + ndim, + poly_order, + data.native, + dims, + weight=weight_native) + + if keep_dirs: + new_grid = [np.asarray(data.grid[d]) for d in keep_dirs] + else: + new_grid = [np.array([0.0, 1.0])] + + return data._result(new_grid, + out_native, + inplace=inplace, + tag=tag, + label=label, + cells=np.asarray(cells_avg)) diff --git a/src/postgkyl/operations/collect.py b/src/postgkyl/operations/collect.py new file mode 100644 index 00000000..f66b4c58 --- /dev/null +++ b/src/postgkyl/operations/collect.py @@ -0,0 +1,121 @@ +"""The ``collect`` verb -- combine many datasets into one along a new time axis.""" + +from __future__ import annotations + +import numpy as np + +from postgkyl.gdatastate import flatten_datasets +from postgkyl.gdatastate.gdatastate import GDataState + + +def _collect_group(states: list, start: int, *, sumdata: bool, + period: float | None, offset: float, tag: str | None, + label: str | None) -> GDataState: + """Collect one group of frames into a single dataset (``collect``'s body, + factored out so ``chunk`` can call it once per chunk). ``start`` is the + group's offset into the full input sequence, so the positional-fallback + time stamp (used when a frame has neither ``ctx['time']`` nor + ``ctx['frame']``) stays consistent with the ungrouped (no-``chunk``) case.""" + time, values = [], [] + grid = None + for i, dat in enumerate(states, start=start): + if dat.backend == "gkyl": + raise ValueError( + f"collect operates on interpolated (NumPy) values; call .interpolate() " + f"first on dataset {i} -- stacking raw DG coefficients would mix " + f"basis functions.") + stamp = dat.ctx.get("time", dat.ctx.get("frame", i)) + time.append(stamp) + + val = dat.values + if sumdata: + values.append(np.nansum(val, axis=tuple(range(dat.num_dims)))) + else: + values.append(val) + if grid is None: + grid = list(dat.grid) + + time = np.array(time) + values = np.array(values) + + if period: + time = (time - offset) % period + + sort_idx = np.argsort(time) + time = time[sort_idx] + values = values[sort_idx] + + out_grid = [time] if sumdata else [np.array(time)] + grid + return states[0]._result(out_grid, + values, + tag=(tag or "default"), + label=(label if label is not None else "collect")) + + +def collect(*datasets: GDataState, + sumdata: bool = False, + period: float | None = None, + offset: float = 0.0, + chunk: int | None = None, + tag: str | None = None, + label: str | None = None) -> GDataState | list[GDataState]: + """Collect many single-frame datasets into one with a new leading time axis. + + Accepts ``collect(a, b)`` or ``collect([a, b])`` (flattened via + ``gdatastate.flatten_datasets``). The per-dataset time stamp is taken from + ``ctx['time']``, then ``ctx['frame']``, then the dataset's position in the + sequence as a fallback; frames are sorted by their (possibly folded) time + stamp. Each result copies the grid/ctx of its group's first frame (via + ``_result``), so it stays the caller's concrete dataset class. + + Args: + *datasets: the datasets to collect (each NumPy-backed, sharing a grid + and component layout), or lists/groups thereof. + sumdata: when True, sum each frame over all of its spatial axes (keeping + components) before stacking, so the output grid is just the time + axis. When False the full spatial data of each frame is retained and + the time axis becomes a new leading dimension. + period: when given, fold the time stamps into one period via + ``(time - offset) % period`` before sorting, producing a phase/epoch + axis instead of an unfolded time axis. + offset: phase offset subtracted before the modulo when ``period`` is + used. + chunk: when given (and non-zero), split the input into consecutive + groups of this length and collect each group separately, returning a + list of datasets (one per chunk; the last chunk may be shorter) + instead of a single dataset. + tag: optional tag for the returned dataset(s). + label: optional label for the returned dataset(s) (defaults to + ``'collect'``). + + Returns: + A single dataset with the collected frames stacked along a new leading + time axis, or, when ``chunk`` is given, a list of such datasets. + + Raises: + ValueError: if there are no datasets to collect, or one is native modal + (gkyl-backed). + """ + states = flatten_datasets(datasets) + if not states: + raise ValueError("collect: no datasets to collect.") + + if chunk: + groups = [(states[i:i + chunk], i) for i in range(0, len(states), chunk)] + return [ + _collect_group(group, + start, + sumdata=sumdata, + period=period, + offset=offset, + tag=tag, + label=label) for group, start in groups + ] + + return _collect_group(states, + 0, + sumdata=sumdata, + period=period, + offset=offset, + tag=tag, + label=label) diff --git a/src/postgkyl/operations/differentiate.py b/src/postgkyl/operations/differentiate.py new file mode 100644 index 00000000..90ee1e4c --- /dev/null +++ b/src/postgkyl/operations/differentiate.py @@ -0,0 +1,104 @@ +"""The ``differentiate`` verb -- numerical gradient of field-domain data. + +Per ``.claude/migration/notes/differentiate-decision.md`` (layer 03): an +*exact* modal derivative would need a ``gpython_basis_eval_grad`` addition to the +compiled shim (``gkeyll/core/zero/gkyl_gpython.h``/``gpython.c`` + +``gpython/csrc/_gpythonmodule.c``), out of scope for every layer above +``gpython``. This +verb instead differentiates *after* ``.interpolate()``, with ``np.gradient`` on the +plain NumPy field values -- a numerical (second-order accurate, cell-centered), not +exact, derivative. Exactness on the modal polynomial is unnecessary here precisely +because the data have already been interpolated to a uniform mesh. + +On a separable axis (the ordinary case, including a nonuniform/stretched +grid), this is a plain per-axis ``np.gradient`` against that axis' own 1-D +coordinate array. On a curvilinear axis -- part of a joint, non-separable +``.map(space="conf")`` block, whose grid arrays are multi-dimensional and +have no single 1-D coordinate of their own -- the physical derivative is +computed via the chain rule instead (``numerics.curvilinear. +physical_gradient``): the whole block's Jacobian is inverted once and reused +for every direction/component request that touches it. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl.numerics import curvilinear + +from ._curvilinear import block_for_axis, curvilinear_blocks + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def differentiate(data: "GDataState", + *, + direction: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Numerical gradient of field-domain data. + + With ``direction=None``, differentiates along every spatial axis and + stacks the results in the component axis (``num_comps`` becomes + ``num_comps * num_dims``, grouped ``[d0_comp0..d0_compN, d1_comp0.., ...]``). + With an explicit ``direction``, differentiates along that one axis only + (``num_comps`` unchanged). A separable axis requires a nodal (edge) grid + one entry longer than the value count along that axis; a mismatched axis + silently returns a wrong result -- a caveat inherited unchanged from the + legacy tool. A curvilinear axis (part of a joint ``.map(space="conf")`` + block) has no such per-axis length convention of its own; its block's + grid arrays carry it instead. + + Args: + data: the dataset to differentiate; must be NumPy-backed (call + ``.interpolate()`` first on native modal data). + direction: 0-based axis to differentiate along; None differentiates + along every axis. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset of the gradient, on ``data``'s (unchanged) grid. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + if data.backend == "gkyl": + raise ValueError( + "differentiate operates on interpolated (NumPy) values; call " + ".interpolate() first -- np.gradient has no basis-space meaning for raw " + "DG coefficients.") + grid = data.grid + values = data.values + num_dims = data.num_dims + nc = values.shape[-1] + + blocks = curvilinear_blocks(grid, data.ctx.get("mapped_axes", {})) + block_grad_cache: dict = {} + + def grad_along(d: int) -> np.ndarray: + info = block_for_axis(blocks, d) + if info is None: + zc = 0.5 * (grid[d][1:] + grid[d][:-1]) # cell centered values + return np.gradient(values, zc, edge_order=2, axis=d) + off, dims = info + if off not in block_grad_cache: + block_coords = [grid[dd] for dd in dims] + block_grad_cache[off] = curvilinear.physical_gradient( + block_coords, values, tuple(dims)) + return block_grad_cache[off][..., dims.index(d)] + + if direction is None: + out_shape = list(values.shape) + out_shape[-1] = nc * num_dims + out_values = np.zeros(out_shape) + for d in range(num_dims): + out_values[..., d * nc:(d + 1) * nc] = grad_along(d) + else: + out_values = grad_along(int(direction)) + return data._result(grid, out_values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/operations/eval_at_coord_proj.py b/src/postgkyl/operations/eval_at_coord_proj.py new file mode 100644 index 00000000..acae0764 --- /dev/null +++ b/src/postgkyl/operations/eval_at_coord_proj.py @@ -0,0 +1,104 @@ +"""The ``eval_at_coord_proj`` verb -- evaluate a native DG field at physical +coordinates in a subset of directions, projecting onto the lower-dimensional +target basis for the survivors, via Gkeyll's ``gkyl_dg_eval_at_coord_proj``. + +Terminal-adjacent (like ``average``): produces a new, lower-dimensional +dataset -- still modal and gkyl-native -- so it composes with further +``.to_nodal()``/``.interpolate()``/``.eval_at_coord_proj()`` calls. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl import dg + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _native_basis(data: "GDataState"): + if data.backend != "gkyl": + raise ValueError( + "eval_at_coord_proj wraps gkyl_dg_eval_at_coord_proj and needs " + "native modal data; it is not available after .interpolate() or " + "without the Gkeyll library.") + if data.ctx.get("value_form", "modal") != "modal": + raise ValueError(f"eval_at_coord_proj expects the modal value_form, not " + f"'{data.ctx['value_form']}'; call .to_modal() first.") + basis_type = data.ctx.get("basis_type") + poly_order = data.ctx.get("poly_order") + if basis_type is None or poly_order is None: + raise ValueError("eval_at_coord_proj: data has no basis_type/poly_order " + "metadata") + return str(basis_type), int(poly_order) + + +def eval_at_coord_proj(data: "GDataState", + eval_dirs, + eval_coords, + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Evaluate ``data`` at ``eval_coords`` in ``eval_dirs`` and project onto + the surviving directions' target basis. + + Args: + data: gkyl-backed (native modal) dataset in the modal value_form. + eval_dirs: 0-based direction indices to evaluate away (e.g. the selected + ``z0``-``z5`` flags at the CLI layer). + eval_coords: physical coordinates, one per entry of ``eval_dirs`` (same + order); in the dataset's own computational grid sense (the same + convention every other native-modal verb here, e.g. ``average``, + uses -- not a separately mapped/deformed physical grid, which this + architecture only ever produces post-``interpolate()``). + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A new dataset over the surviving dimensions -- or a single degenerate + dimension (``grid=[0, 1]``) when every direction is evaluated away, + Gkeyll's own convention since there is no true 0-dimensional basis -- + still modal and gkyl-native. Its basis metadata (``basis_type``, + ``poly_order``, ``num_cdim``, ``num_vdim``) reflects the TARGET basis + Gkeyll picked, which can differ in type from the donor's (e.g. + eliminating a gkhybrid velocity direction can yield a plain serendipity + target). + + Raises: + ValueError: ``data`` is NumPy-backed or non-modal, is missing basis + metadata, or ``eval_dirs``/``eval_coords`` don't match in length or + range. + """ + basis_type, poly_order = _native_basis(data) + ndim = data.num_dims + + grid = { + "ndim": ndim, + "lower": np.asarray(data.ctx["lower"]), + "upper": np.asarray(data.ctx["upper"]), + "cells": np.asarray(data.ctx["cells"]), + } + keep_dirs, cells_tar, out_native, btype_tar, poly_order_tar, cdim_tar, \ + vdim_tar = dg.modal.eval_at_coord_proj(grid, basis_type, ndim, + poly_order, data.native, eval_dirs, eval_coords) + + if keep_dirs: + new_grid = [np.asarray(data.grid[d]) for d in keep_dirs] + else: + new_grid = [np.array([0.0, 1.0])] + + return data._result(new_grid, + out_native, + inplace=inplace, + tag=tag, + label=label, + cells=np.asarray(cells_tar), + basis_type=btype_tar, + poly_order=poly_order_tar, + num_cdim=cdim_tar, + num_vdim=vdim_tar) diff --git a/src/postgkyl/operations/evaluate.py b/src/postgkyl/operations/evaluate.py new file mode 100644 index 00000000..6a9c7f20 --- /dev/null +++ b/src/postgkyl/operations/evaluate.py @@ -0,0 +1,492 @@ +"""The ``evaluate`` verb -- evaluate RPN math expressions over datasets. + +The numeric operators live in :mod:`postgkyl.numerics.ev_ops` (pure +``(grid, values)`` functions, keyed by token in ``numerics.ev_cmds``); this +module is the stack machine that drives them and the glue that resolves +``f``/``fN`` tokens against an explicit list of datasets. + +Expressions use Reverse Polish Notation, e.g. ``"f0 f1 +"`` adds two datasets +and ``"f 2 *"`` doubles one. Data tokens are: + +- ``f`` / ``fN`` -- the ``N``-th provided dataset (``f`` == ``f0``), +- ``fN[c]`` -- component ``c`` of that dataset (slices like ``0:3`` work), +- ``fN.key`` -- the scalar ``ctx[key]`` of that dataset. + +Anything else is parsed as a numeric/axis literal (a float, a ``"0,1"`` / +``"0:3"`` axis spec, or a Python literal in brackets/parens). Every operator +in ``numerics.ev_cmds`` is a plain array function -- none needed a +``NotImplementedError`` GData-only placeholder (see the numerics module +docstring), so there is nothing left to resolve here. + +A data token referencing native (gkyl-backed) data is kept native, not +forced through ``select()``'s point-value guard, regardless of +value_form -- see ``_native_kernel``: + +- **modal** (raw DG coefficients): ``+ - * /`` and integer ``pow``/``sq`` + route through Gkeyll's own weak DG kernels, the same math + ``operations.arithmetic`` uses for the ``GData`` operators. An operator + with no weak-kernel meaning (``sqrt``, ``sin``, reductions, ...) -- or one + Gkeyll's kernel itself refuses for this basis/order -- warns and falls + back to plain NumPy math on the raw coefficient view, rather than + hard-blocking: value_form/basis metadata is sometimes simply wrong (a + diagnostic file mistagged "modal" by its writer; see the load-time + ``--value_form`` override), and the raw view is exact whenever + coefficient 0 already *is* the point value (e.g. p0 data). +- **nodal/quad** (point values): every operator in ``_POINTWISE_TOKENS`` + (``+ - * / pow sq sqrt sin cos tan abs log log10 exp max2 min2 + scale_comp scale_zi_axis``) is exact regardless of packing, so it is + computed with plain NumPy on the view and the result is wrapped back into + a native array -- computed on the view, wrapped back native, staying + in-value_form, mirroring ``operations.arithmetic``'s ufunc dispatch. + Anything else (``dot``, ``avg``, ``max``, ``min``, ``mean``, ``len``, + ``grad``, ``grad2``, ``int``, ``div``, ``curl``) is a genuine reduction or + finite-difference derivative -- not a per-point transform -- so it leaves + the native domain for plain NumPy math on the raw view, same as before; + ``apply_operator`` then strips the now-stale ``value_form`` tag and + marks the result ``interpolated`` (mirroring ``.interpolate()``) so + ``info()`` doesn't keep claiming a value_form the data no longer has. +""" + +from __future__ import annotations + +import re +import warnings +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl import dg +from postgkyl.numerics import ev_cmds +from postgkyl.operations.select import select + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + +# RPN tokens with an exact Gkeyll weak-kernel meaning on modal data. +_MODAL_BINARY_OPS = {"+", "-", "*", "/", "pow"} + +# RPN tokens that are exact, shape-preserving pointwise math on nodal/quad +# point values (elementwise, no cross-cell/cross-node access, no reduction) +# -- safe to compute on the raw view and wrap back into a native array. +# Everything else in numerics.ev_cmds (dot, avg, max, min, mean, len, grad, +# grad2, int, div, curl) is a reduction or a finite-difference derivative +# and must leave the native domain instead. +_POINTWISE_TOKENS = frozenset({ + "+", + "-", + "*", + "/", + "pow", + "sq", + "sqrt", + "sin", + "cos", + "tan", + "abs", + "log", + "log10", + "exp", + "max2", + "min2", + "scale_comp", + "scale_zi_axis", +}) + +# f, f0, f12 ... with optional [comp] selection and optional .ctxkey suffix. +_DATA_TOKEN = re.compile(r"^f(\d*)(?:\[([^\]]*)\])?(?:\.(\w+))?$") + + +def _rep_of(ctx: dict) -> str: + return ctx.get("value_form", "modal") + + +def _compare(a, b) -> bool: + """Equality that also handles NumPy arrays (used when merging ctx dicts).""" + if isinstance(a, np.ndarray): + return np.array_equal(a, b) + return a == b + + +def _modal_view(value, ctx: dict): + """Read-only NumPy view of a native modal operand, for the (warned) + pointwise fallback; anything else passes through unchanged.""" + if dg.modal.is_native(value): + return value.view(ctx.get("cells")) + return value + + +def _basis_of(ctx: dict): + basis_type, poly_order = ctx.get("basis_type"), ctx.get("poly_order") + if basis_type is None or poly_order is None: + raise ValueError("modal operand has no basis_type/poly_order metadata") + return str(basis_type), int(poly_order) + + +def _as_scalar(value): + """A Python float if ``value`` is scalar-shaped, else None.""" + if isinstance(value, (int, float, np.integer, np.floating)): + return float(value) + if isinstance(value, np.ndarray) and value.ndim == 0: + return float(value) + return None + + +def _modal_kernel(token: str, tmp_grid, tmp_values, tmp_ctx): + """Try to compute ``token`` via Gkeyll's own weak DG kernels when a modal + (native, raw-DG-coefficient) operand is present. + + Returns ``(out_grid, out_values)`` when the operator has an exact modal + meaning and Gkeyll's kernel accepts this basis/order (``+ - * /`` and + integer ``pow``/``sq``). Returns ``None`` when no operand is modal + (nothing to do here -- the caller runs the plain NumPy ``func`` as usual). + + Deliberately never raises: basis/value_form metadata can be wrong + (a diagnostic file mistagged "modal" by its writer), so an operator with + no weak-kernel form, a basis/kernel Gkeyll itself refuses, or a + non-scalar second operand all warn and return ``None`` too -- the caller + then falls back to plain NumPy math on the raw coefficient view (exact + whenever coefficient 0 already *is* the point value, e.g. p0 data). + """ + is_modal = [dg.modal.is_native(v) for v in tmp_values] + if not any(is_modal): + return None + + try: + if len(tmp_values) == 1: + if token != "sq": + raise ValueError(f"'{token}' has no weak-kernel form") + basis_type, poly_order = _basis_of(tmp_ctx[0]) + out = dg.modal.power(basis_type, len(tmp_grid[0]), poly_order, + tmp_values[0], 2) + return [tmp_grid[0]], [out] + + if len(tmp_values) == 2: + if token not in _MODAL_BINARY_OPS: + raise ValueError(f"'{token}' has no weak-kernel form") + # RPN order: tmp_values[0] is "b" (top of stack), tmp_values[1] is "a". + a, b = tmp_values[1], tmp_values[0] + a_modal, b_modal = is_modal[1], is_modal[0] + + if a_modal and b_modal: + grid = tmp_grid[1] if tmp_grid[1] is not None else tmp_grid[0] + basis_a, basis_b = _basis_of(tmp_ctx[1]), _basis_of(tmp_ctx[0]) + if basis_a != basis_b: + raise ValueError( + f"operands have different DG bases ({basis_a} vs {basis_b})") + basis_type, poly_order = basis_a + ndim = len(grid) + if token == "+": + out = dg.modal.lincomb(1.0, a, 1.0, b) + elif token == "-": + out = dg.modal.lincomb(1.0, a, -1.0, b) + elif token in ("*", "/"): + fn = dg.modal.weak_mul if token == "*" else dg.modal.weak_div + out = fn(basis_type, ndim, poly_order, a, b) + else: + raise ValueError("'pow' is not defined between two modal datasets") + return [grid], [out] + + # Exactly one operand is modal; the other must be a plain scalar. + modal_arr, modal_ctx, modal_grid = (a, tmp_ctx[1], tmp_grid[1]) if a_modal \ + else (b, tmp_ctx[0], tmp_grid[0]) + other = b if a_modal else a + scalar = _as_scalar(other) + if scalar is None: + raise ValueError("cannot mix native modal data with a plain array") + basis_type, poly_order = _basis_of(modal_ctx) + ndim = len(modal_grid) + scalar_first = not a_modal # the scalar came first in the expression + + if token == "*": + out = dg.modal.scale(modal_arr, scalar) + elif token == "/": + out = (dg.modal.scale( + dg.modal.weak_inv(basis_type, ndim, poly_order, modal_arr), scalar) + if scalar_first else dg.modal.scale(modal_arr, 1.0 / scalar)) + elif token == "+": + out = dg.modal.shift_mean(basis_type, ndim, poly_order, modal_arr, + scalar) + elif token == "-": + out = (dg.modal.shift_mean(basis_type, ndim, poly_order, + dg.modal.scale(modal_arr, -1.0), scalar) + if scalar_first else dg.modal.shift_mean( + basis_type, ndim, poly_order, modal_arr, -scalar)) + else: # pow + if scalar_first or not float(scalar).is_integer() or scalar < 1: + raise ValueError( + f"modal 'pow' needs a modal base and a positive integer " + f"exponent, got exponent {scalar!r} (scalar_first={scalar_first})" + ) + out = dg.modal.power(basis_type, ndim, poly_order, modal_arr, + int(scalar)) + return [modal_grid], [out] + + raise ValueError( + f"'{token}' has no weak-kernel form for {len(tmp_values)} operands") + except Exception as err: + warnings.warn( + f"evaluate: '{token}' on native modal (raw DG coefficient) data: {err}; " + "falling back to plain math on the raw coefficient view -- exact only " + "if coefficient 0 already IS the point value (e.g. p0 data, or a file " + "whose 'modal' tag is wrong; see --value_form).", + stacklevel=3) + return None + + +def _native_kernel(token: str, tmp_grid, tmp_values, tmp_ctx, func): + """Dispatch a native (gkyl-backed) operand to the value_form-correct math. + + Returns ``(out_grid, out_values)`` -- with ``out_values`` wrapped back into + native arrays whenever the result stays a per-point/per-coefficient field + -- or ``None`` when nothing here applies (the caller runs the plain NumPy + ``func`` on the raw view as usual, e.g. for reductions/derivatives). + + - Every native operand modal: delegates to :func:`_modal_kernel` (weak + DG kernels), unchanged. + - Every native operand the *same* nodal/quad value_form, and ``token`` + in :data:`_POINTWISE_TOKENS`: exact NumPy math on the raw view, wrapped + back native -- mirrors ``operations.arithmetic``'s "compute on the view, + wrap back native, stay in-value_form" pointwise dispatch. + - Native operands in *different* value_forms: warns and falls back + (the caller then runs ``func`` on plain views, same as a value_form + mismatch anywhere else in this module). + - Any other token (reductions, finite-difference derivatives): returns + ``None`` so the caller's plain-NumPy path runs -- the result then + genuinely leaves the native/value_form domain. + """ + is_native = [dg.modal.is_native(v) for v in tmp_values] + if not any(is_native): + return None + + reps = { + _rep_of(c) + for v, c, native in zip(tmp_values, tmp_ctx, is_native) if native + } + if reps == {"modal"}: + return _modal_kernel(token, tmp_grid, tmp_values, tmp_ctx) + + if len(reps) > 1: + warnings.warn( + f"evaluate: '{token}' mixes native operands in different " + f"value_forms ({sorted(reps)}); falling back to plain math on " + "the raw views.", + stacklevel=3) + return None + + if token not in _POINTWISE_TOKENS: + return None + + view_values = [_modal_view(v, c) for v, c in zip(tmp_values, tmp_ctx)] + out_grid, out_values = func(tmp_grid, view_values) + return out_grid, [dg.rep.wrap(v) for v in out_values] + + +def apply_operator(grid_stack, value_stack, ctx_stack, token: str) -> bool: + """Reduce the RPN stacks in place by applying ``token`` if it is an operator. + + Each stack entry is a list of "sets" (grids/values/ctx dicts); an operator + pops ``num_in`` entries, applies its pure function from + :data:`postgkyl.numerics.ev_cmds` over every set (broadcasting shorter + inputs), and pushes ``num_out`` results. The ctx of the output is the merge + of the inputs' ctx, dropping any key whose value disagrees between inputs. + + Args: + grid_stack, value_stack, ctx_stack: the parallel RPN stacks, mutated in + place. + token: the candidate operator token (e.g. ``'+'``, ``'sqrt'``, ``'int'``). + + Returns: + True if ``token`` was a known operator and the stacks were reduced; + False if ``token`` is not an operator (the stacks are untouched). + + Raises: + ValueError: if the operator's function raises while evaluating. + """ + if token not in ev_cmds: + return False + num_in = ev_cmds[token]["num_in"] + num_out = ev_cmds[token]["num_out"] + func = ev_cmds[token]["func"] + + in_grid, in_values, in_ctx, num_sets = [], [], [], [] + for _ in range(num_in): + in_grid.append(grid_stack.pop()) + in_values.append(value_stack.pop()) + in_ctx.append(ctx_stack.pop()) + num_sets.append(len(in_values[-1])) + for _ in range(num_out): + grid_stack.append([]) + value_stack.append([]) + ctx_stack.append([]) + + for set_idx in range(max(num_sets)): + tmp_grid, tmp_values, tmp_ctx = [], [], [] + for i in range(num_in): + tmp_grid.append(in_grid[i][min(set_idx, num_sets[i] - 1)]) + tmp_values.append(in_values[i][min(set_idx, num_sets[i] - 1)]) + tmp_ctx.append(in_ctx[i][min(set_idx, num_sets[i] - 1)]) + try: + native_out = _native_kernel(token, tmp_grid, tmp_values, tmp_ctx, func) + if native_out is not None: + out_grid, out_values = native_out + else: + view_values = [_modal_view(v, c) for v, c in zip(tmp_values, tmp_ctx)] + out_grid, out_values = func(tmp_grid, view_values) + except Exception as err: + raise ValueError(str(err)) from err + + # Merge ctx of all inputs; drop keys that disagree between inputs. + out_ctx: dict = {} + remove_list = [] + for i in range(num_in): + for key in tmp_ctx[i]: + if key in out_ctx and _compare(tmp_ctx[i][key], out_ctx[key]): + pass # already copied and matches; nothing to do + elif key in out_ctx: + remove_list.append(key) # discrepancy; mark for removal + else: + out_ctx[key] = tmp_ctx[i][key] + for key in dict.fromkeys(remove_list): + out_ctx.pop(key) + + # A native nodal/quad operand whose result did *not* come back wrapped + # native (a genuine reduction/derivative, per _native_kernel) has left + # the per-point field domain: the merged ctx's 'value_form' is now + # stale (it still names a value_form this output no longer has), so + # drop it and mark the result the same way .interpolate() does -- no + # longer gkyl-native -- rather than let info() keep describing it as a + # value_form it left behind. + was_native_nonmodal = any( + dg.modal.is_native(v) and _rep_of(c) != "modal" + for v, c in zip(tmp_values, tmp_ctx)) + + for i in range(num_out): + grid_stack[-num_out + i].append(out_grid[i]) + value_stack[-num_out + i].append(out_values[i]) + this_ctx = dict(out_ctx) + if was_native_nonmodal and not dg.modal.is_native(out_values[i]): + this_ctx.pop("value_form", None) + this_ctx["interpolated"] = True + ctx_stack[-num_out + i].append(this_ctx) + return True + + +def _push_token(token: str, datasets, grid_stack, value_stack, + ctx_stack) -> bool: + """Push a single non-operator ``token`` (data reference or literal). + + Returns False only if the token cannot be interpreted at all. + """ + match = _DATA_TOKEN.match(token) + if match: + idx = int(match.group(1)) if match.group(1) else 0 + comp = match.group(2) + ctx_key = match.group(3) + dat = datasets[idx] + if ctx_key is not None: + if ctx_key not in dat.ctx: + raise ValueError( + f"evaluate: unknown ctx key '{ctx_key}' on dataset f{idx}") + grid, values = None, np.array(dat.ctx[ctx_key]) + elif comp is None and dat.backend == "gkyl": + # Keep native data on the stack (rather than forcing it through + # select()'s point-value guard), regardless of value_form: RPN + # math routes through Gkeyll's own weak kernels for modal data, or + # exact NumPy math wrapped back native for nodal/quad point values, + # when the operator supports it, or warns/falls back to the raw view + # otherwise -- see _native_kernel. + grid, values = dat.grid, dat.native + else: + # select() carries the shared operability guard (raw modal coefficients + # refuse; nodal/quad value_forms, already point values, pass) for + # a comp-sliced modal token (still genuinely unsafe -- slicing raw DG + # coefficients by component can mix basis functions) and every + # already-point-value token. select() itself now keeps a gkyl-backed + # nodal/quad result native, so keep pushing the native array here too + # (not its plain-view .values) so it stays eligible for _native_kernel. + selected = select(dat, comp=comp) + grid = selected.grid + values = selected.native if selected.backend == "gkyl" else selected.values + grid_stack.append([grid]) + value_stack.append([values]) + ctx_stack.append([dat.ctx]) + return True + + # Numeric / axis literal fallback (mirrors the CLI token parser). + if "(" in token or "[" in token: + value_stack.append([eval(token)]) # noqa: S307 -- trusted expression source + elif ":" in token or "," in token: + value_stack.append([str(token)]) + else: + try: + value_stack.append([np.array(float(token))]) + except ValueError: + return False + grid_stack.append([None]) + ctx_stack.append([{}]) + return True + + +def available_operators() -> list[str]: + """The RPN operator tokens ``evaluate`` recognizes (e.g. ``'+'``, ``'sqrt'``).""" + return sorted(ev_cmds) + + +def evaluate(chain: str, + *datasets: "GDataState", + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Evaluate an RPN expression over an explicit list of datasets. + + ``f``/``fN`` tokens in ``chain`` refer to ``datasets[N]`` (``f`` == ``f0``); + see the module docstring for the token grammar. The result is built via + ``datasets[0]._result(...)`` (so it stays the caller's concrete dataset + class) and holds the single value left on top of the stack. + + Args: + chain: the RPN expression, e.g. ``"f0 f1 +"`` or ``"f sq 2 *"``. + *datasets: the datasets referenced positionally by the ``f``/``fN`` + tokens. At least one is required (it anchors the result's class). + tag: optional tag for the returned dataset (defaults to ``'default'``). + label: optional label for the returned dataset (defaults to ``chain``). + + Returns: + A dataset holding the evaluated grid/values and the merged ctx. + + Raises: + ValueError: if ``datasets`` is empty, the expression is empty, a token + is unrecognized, or an operator fails. + """ + if not datasets: + raise ValueError("evaluate: at least one dataset is required.") + + grid_stack, value_stack, ctx_stack = [], [], [] + for token in filter(None, chain.split(" ")): + if apply_operator(grid_stack, value_stack, ctx_stack, token): + continue + if not _push_token(token, datasets, grid_stack, value_stack, ctx_stack): + raise ValueError( + f"evaluate: token '{token}' is neither data nor an operator") + + if not value_stack: + raise ValueError("evaluate: expression produced no result") + + final_grid = grid_stack[-1][0] + final_values = value_stack[-1][0] + final_ctx = dict(ctx_stack[-1][0]) + out_grid = final_grid if final_grid is not None else datasets[0].grid + result = datasets[0]._result(out_grid, + final_values, + tag=(tag or "default"), + label=(label if label is not None else chain)) + # The result's ctx is the RPN merge (apply_operator already resolved every + # conflict), not datasets[0]'s ctx that '_result' copied as a starting + # point -- a key apply_operator dropped as conflicting must not survive + # just because it happened to be on datasets[0]. 'cells'/'num_comps'/ + # 'lower'/'upper' are the shape/grid-derived facts '_result's push() just + # recomputed from the actual final_grid/final_values; keep those. + derived = {"cells", "num_comps", "lower", "upper"} + kept = {k: result.ctx[k] for k in derived if k in result.ctx} + result.ctx = final_ctx + result.ctx.update(kept) + return result diff --git a/src/postgkyl/operations/extract_input.py b/src/postgkyl/operations/extract_input.py new file mode 100644 index 00000000..cea179b1 --- /dev/null +++ b/src/postgkyl/operations/extract_input.py @@ -0,0 +1,34 @@ +"""The ``extract_input`` verb -- decode the input file embedded in ``ctx``. + +Gkeyll output files may carry the original simulation input file as a +base64-encoded string, stashed by the reader under ``ctx['input_file']``. +This verb is *terminal*: unlike every other verb in this module it returns +a plain ``str``, not a dataset (matching the legacy contract). + +No current :mod:`postgkyl.io` reader populates ``ctx['input_file']``; this +verb decodes it whenever a reader does provide it, and returns ``""`` +otherwise, exactly as the legacy code did when no input file was embedded. +""" + +from __future__ import annotations + +import base64 +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def extract_input(data: "GDataState") -> str: + """Decode the input file embedded in a Gkeyll output file's ``ctx``. + + Args: + data: the dataset whose embedded input file is decoded. + + Returns: + The decoded input-file text, or an empty string when none is embedded. + """ + encoded = data.ctx.get("input_file") + if encoded: + return base64.decodebytes(encoded.encode("utf-8")).decode("utf-8") + return "" diff --git a/src/postgkyl/operations/fft.py b/src/postgkyl/operations/fft.py new file mode 100644 index 00000000..980e4329 --- /dev/null +++ b/src/postgkyl/operations/fft.py @@ -0,0 +1,60 @@ +"""The ``fft`` verb -- Fourier transform / power spectral density.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import numerics + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def fft(data: "GDataState", + *, + psd: bool = False, + iso: bool = False, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Fourier transform (or power spectral density) of field-domain data. + + Wraps ``numerics.fft``: each component is transformed over the spatial + axes (dummy axes of length <= 2 are squeezed out first). Supports 1D, 2D, + and 3D data. ``numerics.fft`` reads its sample spacing straight off the + grid array's own length, so a nodal (edge) grid -- one entry longer than + the value count, the usual post-``.interpolate()`` shape -- is first collapsed + to cell centers (matching values); a grid that already matches (e.g. a + dynvector's) is passed through unchanged. + + Args: + data: the dataset to transform; must be NumPy-backed (call ``.interpolate()`` + first on native modal data). + psd: when True, return the power spectral density ``|FT|^2`` over the + positive frequencies only. + iso: when True (only meaningful for 2D/3D data with ``psd=True``), bin + the PSD into a 1D isotropic spectrum over the polar wavenumber + magnitude. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset whose grid is the frequency/wavenumber axis (axes) and whose + values are the transform, PSD, or isotropic spectrum. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or if isotropic + binning is requested for data that is not 2D/3D. + """ + if data.backend == "gkyl": + raise ValueError( + "fft operates on interpolated (NumPy) values; call .interpolate() first " + "-- Fourier transforming raw DG coefficients would mix basis functions." + ) + grid, values = data.grid, data.values + num_cells = values.shape[:-1] + if any(grid[d].shape[0] == num_cells[d] + 1 for d in range(len(grid))): + grid = numerics.nodal_to_cell_centered_grid(grid, num_cells) + freq, ft_values = numerics.fft(grid, values, psd=psd, iso=iso) + return data._result(freq, ft_values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/operations/fit.py b/src/postgkyl/operations/fit.py new file mode 100644 index 00000000..a2f89095 --- /dev/null +++ b/src/postgkyl/operations/fit.py @@ -0,0 +1,212 @@ +"""The ``fit`` verb -- fit a model to data and return the fitted curve. + +The result holds the fitted values on the data's grid; the per-component fit +parameters, 1-sigma uncertainties, and R^2 are stored in +``ctx['fit_params']``, ``ctx['fit_std']``, and ``ctx['fit_R2']``. ``fit_type`` +is a model name (e.g. ``'linear'``, ``'gaussian'``, ``'exp2'`` for +growth-rate fits) or an RPN expression -- see :mod:`postgkyl.numerics.fit`. + +``window=True`` restricts each component's fit to its best-scoring leading +window rather than the full domain -- the growth-rate use case, where only +a continuously growing/decaying leading region of a longer time series +should be fit (e.g. ``fit(d, 'exp2', window=True)``); see +:func:`postgkyl.numerics.fit_best_window`. +""" + +from __future__ import annotations + +import inspect +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl import numerics + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def fit(data: "GDataState", + fit_type: str, + *, + guess=None, + window: bool = False, + min_n: int | None = None, + print_coeffs: bool = False, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Fit a model to data and return the fitted curve. + + Fits the model named (or expressed) by ``fit_type`` to each component of + ``data`` independently and returns the fitted values evaluated on the + data's (cell-centered) grid. Axes collapsed to a single cell (e.g. after + ``integrate`` or ``select``) are dropped, so 1D and 2D fits are supported. + + Args: + data: the dataset to fit; must be NumPy-backed. Its grid provides the + independent variable(s) and each component is fit separately. + fit_type: the model to fit -- a key of ``numerics.FIT_FUNCTIONS`` + ('linear', 'quadratic', 'plane', 'quadratic2d', 'exp_plateau', + 'gaussian', 'power', 'sinusoid', 'tanh_transition', 'exp2'), or a + custom RPN expression string (e.g. ``'x a * b +'``) whose free tokens + (not the spatial variables 'x'/'y', operators, or numbers) become fit + parameters. + guess: initial guess for the fit parameters -- a comma-separated string + (e.g. ``'1,0,2'``) or a sequence of floats. None derives a + data-driven guess per component via ``numerics.auto_guess`` (for the + first window, if ``window=True``). + window: fit only the best-scoring leading window of the data (1D only) + instead of the full domain -- see ``numerics.fit_best_window``. + min_n: minimum window length when ``window=True``; ``None`` defaults to + one tenth of the number of samples. Ignored otherwise. + print_coeffs: print the model equation and coefficient descriptions, + followed by named coefficients with estimated 1-sigma uncertainties, + R^2, residual sum of squares (RSS), root mean squared error (RMSE), + residual standard error, sample count, degrees of freedom, and + coordinate ranges for each zero-based component (12 significant + digits). Statistics use only the fitted points, including when + ``window=True``. R^2 is undefined for constant data; residual standard + error requires positive residual degrees of freedom. Custom RPN models + show their expression and free parameter names. Full precision remains + in ``ctx['fit_params']``, ``ctx['fit_std']``, and ``ctx['fit_R2']``. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset holding the fitted curve on the active grid, with + ``ctx['fit_params']``, ``ctx['fit_std']``, and ``ctx['fit_R2']`` set. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), if ``fit_type`` + is neither a recognized model name nor a valid RPN expression, if the + data's active dimensionality does not match the model's, or if + ``window=True`` and the data is not 1D. + """ + if data.backend == "gkyl": + raise ValueError( + "fit operates on interpolated (NumPy) values; call .interpolate() first " + "-- fitting raw DG coefficients would mix basis functions.") + grid = data.grid + values = data.values + spatial_shape = values.shape[:-1] + + if any(grid[d].shape[0] == spatial_shape[d] + 1 for d in range(len(grid))): + cc_grid = numerics.nodal_to_cell_centered_grid(grid, spatial_shape) + else: + cc_grid = list(grid) + + # Drop dimensions collapsed to a single cell (e.g. after integrate/select). + active = [d for d in range(len(cc_grid)) if cc_grid[d].shape[0] > 1] + if len(active) < len(cc_grid): + idx = tuple( + slice(None) if d in active else 0 + for d in range(len(spatial_shape))) + (slice(None), ) + cc_grid = [cc_grid[d] for d in active] + values = values[idx] + + ndim_fit = numerics.FIT_NDIM.get(fit_type, numerics.rpn_ndim(fit_type)) + if len(cc_grid) != ndim_fit: + raise ValueError( + f"fit '{fit_type}' requires {ndim_fit:d} spatial dimension(s), but " + f"data has {len(cc_grid):d}. Reduce it first (e.g. select or integrate)." + ) + if window and len(cc_grid) != 1: + raise ValueError( + "fit: window=True is only supported for 1D (time-series-like) data, " + f"but data has {len(cc_grid):d} active dimension(s).") + + if len(cc_grid) == 1: + xdata = cc_grid[0] + else: + mesh = np.meshgrid(cc_grid[0], cc_grid[1], indexing="ij") + xdata = np.array([mesh[0].flatten(), mesh[1].flatten()]) + + guess_list = None + if guess is not None: + guess_list = ([float(v) for v in guess.split(",")] if isinstance( + guess, str) else list(guess)) + + active_shape = tuple(cg.shape[0] for cg in cc_grid) + fit_values_list, all_params, all_std, all_r2 = [], [], [], [] + all_n = [] + for comp in range(values.shape[-1]): + ydata = values[..., comp].flatten() + if window: + params, cov, r2, n = numerics.fit_best_window(xdata, + ydata, + fit_type, + min_n=min_n, + p0=guess_list) + else: + n = ydata.size + p0 = guess_list if guess_list is not None else numerics.auto_guess( + fit_type, xdata, ydata) + params, cov, r2 = numerics.fit(xdata, ydata, fit_type, p0=p0) + y_fit = numerics.fit_evaluate(xdata, fit_type, params) + fit_values_list.append(y_fit.reshape(active_shape + (1, ))) + all_params.append(params) + all_std.append(np.sqrt(np.diag(cov))) + all_r2.append(r2) + all_n.append(n) + + fit_values = np.concatenate(fit_values_list, axis=-1) + fit_grid = [grid[d] for d in active] + if print_coeffs: + model = numerics.FIT_FUNCTIONS.get(fit_type) + if model is None: + description = f"RPN expression: {fit_type}" + param_names = numerics.rpn_param_names(fit_type) + else: + description = inspect.getdoc(model).replace("``", "") + param_names = list(inspect.signature(model).parameters)[1:] + print(f"fit '{fit_type}':") + for line in description.splitlines(): + if line: + print(f" {line}") + if len(active) == 1: + print(" x: input coordinate (time for a time series).") + else: + print(f" x, y: input coordinates on grid axes {active[0]}, {active[1]}.") + for comp, params in enumerate(all_params): + print(f" component {comp}:") + for name, value, std in zip(param_names, params, all_std[comp]): + if np.isfinite(std): + print(f" {name} = {value:.12g} +/- {std:.12g} (1-sigma)") + else: + print(f" {name} = {value:.12g} (1-sigma uncertainty unavailable)") + n = all_n[comp] + observed = values[..., comp].reshape(-1)[:n] + predicted = fit_values[..., comp].reshape(-1)[:n] + rss = np.sum((observed - predicted)**2) + dof = n - len(params) + sample_scope = (f" of {values[..., comp].size} (leading window)" + if window else "") + print(f" Samples = {n}{sample_scope}") + print(f" Parameters = {len(params)}; " + f"residual degrees of freedom = {dof}") + for name, coordinates in zip(("x", "y"), cc_grid): + fitted_coordinates = coordinates[:n] if window else coordinates + print(f" {name} range = [{fitted_coordinates.min():.12g}, " + f"{fitted_coordinates.max():.12g}]") + if np.any(observed != observed[0]): + print(f" R^2 = {all_r2[comp]:.12g} (coefficient of determination)") + else: + print(" R^2 = undefined (constant fitted data)") + print(f" RSS = {rss:.12g} (sum of squared residuals)") + print(f" RMSE = {np.sqrt(rss / n):.12g} (root mean squared error)") + if dof > 0: + print(f" Residual standard error = {np.sqrt(rss / dof):.12g} " + "(sqrt(RSS / degrees of freedom))") + else: + print(" Residual standard error = undefined " + "(no residual degrees of freedom)") + return data._result(fit_grid, + fit_values, + inplace=inplace, + tag=tag, + label=label, + fit_params=all_params, + fit_std=all_std, + fit_R2=all_r2) diff --git a/src/postgkyl/operations/grid.py b/src/postgkyl/operations/grid.py new file mode 100644 index 00000000..110ddc38 --- /dev/null +++ b/src/postgkyl/operations/grid.py @@ -0,0 +1,64 @@ +"""The ``grid`` verb -- turn a dataset's grid into a dataset of coordinates.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def grid(data: "GDataState", + *, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Turn a dataset's grid into a dataset of coordinate values. + + Builds a new dataset whose values, at each grid node, are the physical + coordinates of ``data``'s grid (one component per dimension). Handles + uniform meshes, separable (velocity) mappings, and full curvilinear mapped + grids (produced by the ``map`` verb) alike. + + Args: + data: the dataset whose grid is converted to coordinate values; must be + NumPy-backed. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset with one component per dimension holding the physical + coordinates, on a placeholder index grid (one cell per original node). + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or its grid does + not have one entry per dimension reported by ``num_cells``. + """ + if data.backend == "gkyl": + raise ValueError( + "grid operates on interpolated (NumPy) values; call .interpolate() " + "first -- raw DG coefficients have no per-node coordinates.") + grid_in = data.grid + num_dims = data.num_dims + num_cells = data.num_cells + if len(grid_in) != num_dims: + raise ValueError( + f"grid: dataset reports {num_dims:d} dimension(s) but its grid has " + f"{len(grid_in):d} axis (axes); shapes are inconsistent.") + + grid_out = [np.arange(nc + 2) for nc in num_cells] + + shape = np.append(np.copy(num_cells) + 1, num_dims) + values = np.zeros(shape) + if num_dims == 1: + values[..., 0] = grid_in[0] + elif len(grid_in[0].shape) == 1: # uniform mesh or separable mapping + for d, t in enumerate(np.meshgrid(*grid_in, indexing="ij")): + values[..., d] = t + else: # curvilinear mapped grid + for d, t in enumerate(grid_in): + values[..., d] = t + return data._result(grid_out, values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/operations/growth.py b/src/postgkyl/operations/growth.py new file mode 100644 index 00000000..9a60fd54 --- /dev/null +++ b/src/postgkyl/operations/growth.py @@ -0,0 +1,40 @@ +"""Convenience composition for exponential growth-rate fits.""" + +from __future__ import annotations + +from postgkyl.gdatastate.gdatastate import GDataState + +from .fit import fit + + +def growth(data: GDataState, + *, + guess: str | None = None, + min_n: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> GDataState: + """Fit exponential growth on the best leading data window. + + Args: + data: One-dimensional point-value dataset to fit. + guess: Initial ``amplitude,rate`` parameter guess. + min_n: Minimum number of points in the fitted leading window. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the fitted curve. + label: Optional label for the fitted curve. + + Returns: + The fitted curve; rate, uncertainty, and R-squared are in its fit context. + """ + return fit(data, + "exp2", + guess=guess, + window=True, + min_n=min_n, + inplace=inplace, + tag=tag, + label=label) + + +__all__ = ["growth"] diff --git a/src/postgkyl/operations/gyrokinetics/__init__.py b/src/postgkyl/operations/gyrokinetics/__init__.py new file mode 100644 index 00000000..de479aba --- /dev/null +++ b/src/postgkyl/operations/gyrokinetics/__init__.py @@ -0,0 +1,47 @@ +"""Gyrokinetic data transformations. + +Placement answers two independent questions: ``operations`` says these +functions re-express data rather than derive physical conclusions, while +``gyrokinetics`` identifies the domain knowledge their geometry requires. +""" + +from .geometry import GKYL_GEOMETRY_ID, Geometry, is_geo_mapc2p, resolve_geometry +from .rz import RzProjection, gk_rz, map_to_rz, resolve_rz_projection +from .fluxsurf import ( + FluxSurfaceGrid, + extract_flux_surface, + gk_fluxsurf, + resolve_flux_surface_grid, +) + +from postgkyl.cli_spec import CommandSpec, Execution, Section, command, hidden +from postgkyl.gdatastate.gdatastate import GDataState + +gk_rz.__globals__.setdefault("GDataState", GDataState) +gk_fluxsurf.__globals__.setdefault("GDataState", GDataState) +command(CommandSpec(Section.VERBS, Execution.MAP_REPLACE))(gk_rz) +command(CommandSpec(Section.VERBS, Execution.MAP_REPLACE))(gk_fluxsurf) +for _function in ( + is_geo_mapc2p, + resolve_geometry, + map_to_rz, + resolve_rz_projection, + extract_flux_surface, + resolve_flux_surface_grid, +): + hidden("lower-level geometry API requires Python geometry objects")(_function) + +__all__ = [ + "GKYL_GEOMETRY_ID", + "Geometry", + "is_geo_mapc2p", + "resolve_geometry", + "RzProjection", + "gk_rz", + "map_to_rz", + "resolve_rz_projection", + "FluxSurfaceGrid", + "extract_flux_surface", + "gk_fluxsurf", + "resolve_flux_surface_grid", +] diff --git a/src/postgkyl/operations/gyrokinetics/fluxsurf.py b/src/postgkyl/operations/gyrokinetics/fluxsurf.py new file mode 100644 index 00000000..0bb203c2 --- /dev/null +++ b/src/postgkyl/operations/gyrokinetics/fluxsurf.py @@ -0,0 +1,218 @@ +"""Extract theta-phi flux surfaces from gyrokinetic field-aligned data.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from scipy.interpolate import PchipInterpolator + +from .geometry import ( + Geometry, + geometry_prefix, + _interpolate_component, + _interpolation_grid, + _resample_grid, + _same_grid, + _validate_component, + _validate_geometry, + _validate_modal_data, + _validate_positive_int, + per_block_path, + resolve_geometry, +) + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +@dataclass(frozen=True) +class FluxSurfaceGrid: + """Precomputed toroidal sampling grid for one radial flux surface.""" + + x_idx: int + zc: np.ndarray + zf: np.ndarray + phi_tor_list: np.ndarray + phi_2d: np.ndarray + computational_grid: tuple[np.ndarray, ...] | None = None + + +def resolve_flux_surface_grid(first: "GDataState", + geo: Geometry, + *, + x_idx: int = 0, + nphi: int = 128, + nz_interp: int = 8) -> FluxSurfaceGrid: + """Precompute a flux-surface sampling grid for compatible 3-D fields.""" + _validate_modal_data(first, "gk_fluxsurf", (3, )) + nphi = _validate_positive_int(nphi, "nphi") + nz_interp = _validate_positive_int(nz_interp, "nz_interp") + _validate_geometry(geo, 3) + if geo.phi is None: + raise ValueError( + "The geometry file has no toroidal-angle component; cannot extract a flux surface." + ) + if isinstance(x_idx, bool) or not isinstance(x_idx, (int, np.integer)): + raise ValueError("x_idx must be an integer radial index.") + + edges, centers = _interpolation_grid(first) + xc, yc, zc = centers + x_idx = int(x_idx) + if not 0 <= x_idx < xc.size: + raise ValueError( + f"x_idx {x_idx} is out of bounds for data with Nx={xc.size}.") + if zc.size < 2 or yc.size < 2: + raise ValueError( + "gk_fluxsurf requires at least two interpolated y and z points.") + + zf_edges = np.linspace(edges[2][0], edges[2][-1], nz_interp * zc.size + 1) + zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) + phi = np.unwrap(np.unwrap(np.unwrap(geo.phi, axis=2), axis=1), axis=0) + phi_grid = _resample_grid(phi, geo.coords, [xc, yc, zf]) + phi_2d = phi_grid[x_idx, :, :] + phi_tor_list = np.linspace(0.0, 2.0 * np.pi, nphi, endpoint=False) + return FluxSurfaceGrid(x_idx=x_idx, + zc=zc, + zf=zf, + phi_tor_list=phi_tor_list, + phi_2d=phi_2d, + computational_grid=tuple( + np.array(axis, copy=True) for axis in edges)) + + +def extract_flux_surface(data: "GDataState", + fs_grid: FluxSurfaceGrid, + *, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Extract component ``comp`` using a reusable ``fs_grid``.""" + _validate_modal_data(data, "gk_fluxsurf", (3, )) + _validate_component(data, comp) + edges, _ = _interpolation_grid(data) + if fs_grid.computational_grid is not None \ + and not _same_grid(fs_grid.computational_grid, edges): + raise ValueError( + "Incompatible flux-surface grid: data computational grid does not " + "match the grid used to build the projection.") + nx, ny, nz = (axis.size - 1 for axis in edges) + if not 0 <= fs_grid.x_idx < nx: + raise ValueError( + f"x_idx {fs_grid.x_idx} is out of bounds for data with Nx={nx}.") + if (fs_grid.zc.shape != (nz, ) + or fs_grid.phi_2d.shape != (ny, fs_grid.zf.size) + or fs_grid.phi_tor_list.ndim != 1): + raise ValueError( + "Incompatible flux-surface grid: projection and data grid shapes differ." + ) + + _, _, values = _interpolate_component(data, comp) + vals_zf = PchipInterpolator(fs_grid.zc, values, axis=-1, + extrapolate=True)(fs_grid.zf) + vals_2d = vals_zf[fs_grid.x_idx, :, :] + + flux_surf_data = np.empty((fs_grid.phi_tor_list.size, fs_grid.zf.size)) + for iz in range(fs_grid.zf.size): + phi_y = fs_grid.phi_2d[:, iz] + val_y = vals_2d[:, iz] + box = np.mean(np.diff(phi_y)) * ny + if not np.isfinite(box) or np.isclose(box, 0.0): + raise ValueError( + "Toroidal geometry has a zero or non-finite binormal angular span.") + phi_ext = np.concatenate([phi_y - box, phi_y, phi_y + box]) + val_ext = np.concatenate([val_y, val_y, val_y]) + order = np.argsort(phi_ext) + folded = phi_y[0] + np.mod(fs_grid.phi_tor_list - phi_y[0], box) + flux_surf_data[:, iz] = np.interp(folded, phi_ext[order], val_ext[order]) + + return data._result([fs_grid.phi_tor_list, fs_grid.zf], + flux_surf_data[..., np.newaxis], + inplace=inplace, + tag=tag, + label=label, + interpolated=True) + + +def flux_surface_grids(datasets, + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + x_idx: int = 0, + nphi: int = 128, + nz_interp: int = 8) -> dict[str | None, FluxSurfaceGrid]: + """Build one reusable flux-surface grid per block geometry.""" + grids: dict[str | None, FluxSurfaceGrid] = {} + for data in datasets: + key = geometry_prefix(data.file_name) + if key in grids: + continue + block = data.ctx.get("block") + geometry = resolve_geometry(data.file_name, + mapc2p=per_block_path(mapc2p, block), + nodes_file=per_block_path(nodes_file, block)) + grids[key] = resolve_flux_surface_grid(data, + geometry, + x_idx=x_idx, + nphi=nphi, + nz_interp=nz_interp) + return grids + + +def grid_for(grids: dict[str | None, FluxSurfaceGrid], + data: "GDataState") -> FluxSurfaceGrid: + """Return the sampling grid belonging to ``data``'s block.""" + return grids[geometry_prefix(data.file_name)] + + +def gk_fluxsurf(data: "GDataState", + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + x_idx: int = 0, + nphi: int = 128, + nz_interp: int = 8, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Extract one field component on a toroidal flux surface. + + Args: + data: Three-dimensional field-aligned modal dataset. + mapc2p: Explicit modal geometry path. + nodes_file: Explicit nodal geometry path. + x_idx: Radial cell index identifying the surface. + nphi: Number of toroidal-angle slices. + nz_interp: Parallel-direction interpolation factor. + comp: Physical field component to extract. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + geometry = resolve_geometry(data.file_name, + mapc2p=mapc2p, + nodes_file=nodes_file) + fs_grid = resolve_flux_surface_grid(data, + geometry, + x_idx=x_idx, + nphi=nphi, + nz_interp=nz_interp) + return extract_flux_surface(data, + fs_grid, + comp=comp, + inplace=inplace, + tag=tag, + label=label) + + +__all__ = [ + "FluxSurfaceGrid", + "extract_flux_surface", + "flux_surface_grids", + "gk_fluxsurf", + "grid_for", + "resolve_flux_surface_grid", +] diff --git a/src/postgkyl/operations/gyrokinetics/geometry.py b/src/postgkyl/operations/gyrokinetics/geometry.py new file mode 100644 index 00000000..c1c99fc0 --- /dev/null +++ b/src/postgkyl/operations/gyrokinetics/geometry.py @@ -0,0 +1,355 @@ +"""Shared geometry machinery for gyrokinetic data transformations. + +This module owns geometry-file discovery and loading plus the grid helpers +used by both the R-Z and flux-surface operations. It deliberately constructs +the verb-less :class:`~postgkyl.gdatastate.gdatastate.GDataState` and calls the +lower interpolation operation directly; the operation layer never reaches up +through the fluent :class:`postgkyl.gdata.GData` surface. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +import numpy as np +from scipy.interpolate import RegularGridInterpolator + +from postgkyl.dg import num_basis +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.io import parse_output_name +from postgkyl.numerics import nodal_to_cell_centered_grid + +from ..interpolate import interpolate + +# Mirrors ``enum gkyl_geometry_id`` in gkeyll/core/zero/gkyl_eqn_type.h. +# This foreign-format fact is shared with the grid-node diagnostic, which +# imports it from here instead of maintaining a second copy. +GKYL_GEOMETRY_ID = [ + "GKYL_GEOMETRY_NONE", + "GKYL_GEOMETRY_TOKAMAK", + "GKYL_GEOMETRY_MIRROR", + "GKYL_GEOMETRY_MAPC2P", + "GKYL_GEOMETRY_FROMFILE", +] +_MAPC2P_IDX = GKYL_GEOMETRY_ID.index("GKYL_GEOMETRY_MAPC2P") + + +@dataclass(frozen=True) +class Geometry: + """Physical ``(R, Z[, phi])`` geometry on its own point grid. + + ``corner`` closes the poloidal domain's ``theta = +/-pi`` ends for a 3-D + R-Z projection. It is ``None`` when no + ``'-geo_corn_nodes.gkyl'`` file exists alongside the field. + """ + + coords: list[np.ndarray] + major_r: np.ndarray + vert_z: np.ndarray + phi: np.ndarray | None + corner: tuple[list[np.ndarray], np.ndarray, np.ndarray] | None + + +def is_geo_mapc2p(ctx: dict) -> bool: + """Whether ``ctx`` identifies user-supplied Cartesian MAPC2P geometry. + + Files without ``geometry_type`` retain the historical MAPC2P default. + """ + return ctx.get("geometry_type", _MAPC2P_IDX) == _MAPC2P_IDX + + +def geometry_prefix(file_name: str | None) -> str | None: + """Return the per-block simulation prefix for ``file_name``. + + Parsing is delegated to :mod:`postgkyl.io.naming`, the authoritative home + of Gkeyll's output-name convention. + """ + name = parse_output_name(file_name) + return name.prefix if name is not None else None + + +def per_block_path(path: str | None, block: int | None) -> str | None: + """Substitute a multiblock index for ``'*'`` in a geometry override.""" + if path is None or block is None or "*" not in path: + return path + return path.replace("*", str(block)) + + +def _gauss_nodes(edges: np.ndarray) -> np.ndarray: + """Physical p1 Gauss-node coordinates for a one-dimensional edge grid.""" + centers = 0.5 * (edges[:-1] + edges[1:]) + offsets = np.diff(edges) / (2.0 * np.sqrt(3.0)) + return np.ravel(np.column_stack([centers - offsets, centers + offsets])) + + +def _pointwise_file( + path: str) -> tuple[list[np.ndarray], np.ndarray, GDataState]: + """Read a point-value geometry file and squeeze singleton dimensions.""" + data = GDataState(path) + grid = [np.squeeze(axis) for axis in data.grid] + return grid, np.squeeze(data.values), data + + +def _geometry_components( + values: np.ndarray, data: GDataState, + path: str) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: + """Interpret one geometry value array as ``R``, ``Z``, and optional phi.""" + required = 3 if is_geo_mapc2p(data.ctx) else 2 + if values.ndim < 2 or values.shape[-1] < required: + kind = "Cartesian X/Y/Z" if required == 3 else "R/Z" + raise ValueError( + f"Geometry file '{path}' must contain at least {required} {kind} components." + ) + + if is_geo_mapc2p(data.ctx): + x, y, z = values[..., 0], values[..., 1], values[..., 2] + return np.sqrt(x**2 + y**2), z, np.arctan2(y, x) + r, z = values[..., 0], values[..., 1] + phi = values[..., 2] if r.ndim == 3 and values.shape[-1] >= 3 else None + return r, z, phi + + +def _read_mapc2p_geometry(path: str): + """Interpolate a modal geometry file to physical ``R``, ``Z``, and phi.""" + source = GDataState(path) + field = interpolate(source) + cells = field.values.shape[:-1] + coords = nodal_to_cell_centered_grid(field.grid, cells) + major_r, vert_z, phi = _geometry_components(field.values, field, path) + return coords, major_r, vert_z, phi + + +def _read_nodes_geometry(path: str): + """Read a p1 pointwise nodal geometry file.""" + grid, values, data = _pointwise_file(path) + coords = [] + for dim, axis in enumerate(grid): + if axis.ndim != 1 or axis.shape[0] != values.shape[dim] + 1 \ + or values.shape[dim] % 2: + raise ValueError(f"Unrecognized nodal geometry layout in '{path}'.") + coords.append(_gauss_nodes(axis[::2])) + major_r, vert_z, phi = _geometry_components(values, data, path) + return coords, major_r, vert_z, phi + + +def _read_corner_rz(path: str): + """Read R and Z from a pointwise ``'-geo_corn_nodes.gkyl'`` file.""" + grid, values, data = _pointwise_file(path) + coords = [ + np.linspace(axis[0], axis[-1], n) + for axis, n in zip(grid, values.shape[:-1]) + ] + major_r, vert_z, _ = _geometry_components(values, data, path) + return coords, major_r, vert_z + + +def _validate_geometry(geometry: Geometry, num_dims: int) -> None: + """Validate geometry tensor shapes for a data grid of ``num_dims``.""" + if len(geometry.coords) != num_dims: + raise ValueError( + f"Geometry has {len(geometry.coords)} dimensions but the data is " + f"{num_dims}-D.") + if any( + np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 + for axis in geometry.coords): + raise ValueError( + "Geometry coordinates must be one-dimensional arrays with at least two points." + ) + if any(not (np.all(np.diff(axis) > 0) or np.all(np.diff(axis) < 0)) + for axis in geometry.coords): + raise ValueError("Geometry coordinate arrays must be strictly monotonic.") + shape = tuple(np.asarray(axis).size for axis in geometry.coords) + if geometry.major_r.shape != shape or geometry.vert_z.shape != shape: + raise ValueError( + "Geometry coordinate and R/Z array shapes are incompatible: " + f"expected {shape}, got R{geometry.major_r.shape} and Z{geometry.vert_z.shape}." + ) + if geometry.phi is not None and geometry.phi.shape != shape: + raise ValueError( + f"Geometry toroidal-angle shape {geometry.phi.shape} does not match {shape}." + ) + if geometry.corner is not None: + corner_coords, corner_r, corner_z = geometry.corner + if len(corner_coords) != num_dims: + raise ValueError( + f"Corner geometry has {len(corner_coords)} dimensions; expected {num_dims}." + ) + corner_shape = tuple(np.asarray(axis).size for axis in corner_coords) + if (any( + np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 + for axis in corner_coords) or corner_r.shape != corner_shape + or corner_z.shape != corner_shape): + raise ValueError( + "Corner geometry coordinate and R/Z array shapes are incompatible.") + + +def resolve_geometry(file_name: str | None, + *, + mapc2p: str | None = None, + nodes_file: str | None = None) -> Geometry: + """Resolve and load the geometry belonging to ``file_name``. + + The exact pointwise ``'-geo_int_nodes.gkyl'`` representation is + preferred, with ``'-geo_int_mapc2p.gkyl'`` as the modal fallback. + ``nodes_file`` and ``mapc2p`` override that lookup and are mutually + exclusive. Passing ``mapc2p=''`` explicitly requests the inferred modal + filename. + + Raises: + ValueError: If both overrides are supplied or no geometry can be found. + """ + if mapc2p is not None and nodes_file is not None: + raise ValueError("Pass either mapc2p= or nodes_file=, not both.") + + parsed = parse_output_name(file_name) + prefix = geometry_prefix(file_name) + block = parsed.block if parsed is not None else None + nodes_file = per_block_path(nodes_file, block) + mapc2p = per_block_path(mapc2p, block) + if nodes_file is not None: + path, kind = nodes_file, "nodes" + elif mapc2p is not None: + path = mapc2p or (f"{prefix}-geo_int_mapc2p.gkyl" if prefix else None) + kind = "mapc2p" + elif prefix is not None: + path, kind = f"{prefix}-geo_int_nodes.gkyl", "nodes" + if not os.path.exists(path): + path, kind = f"{prefix}-geo_int_mapc2p.gkyl", "mapc2p" + else: + path, kind = None, None + + if path is None or not os.path.exists(path): + raise ValueError( + "Could not find a geometry file; pass nodes_file= or mapc2p= explicitly." + ) + + coords, major_r, vert_z, phi = (_read_nodes_geometry(path) if kind == "nodes" + else _read_mapc2p_geometry(path)) + + corner = None + if prefix is not None: + corner_path = f"{prefix}-geo_corn_nodes.gkyl" + if os.path.exists(corner_path): + corner = _read_corner_rz(corner_path) + + geometry = Geometry(coords=coords, + major_r=major_r, + vert_z=vert_z, + phi=phi, + corner=corner) + _validate_geometry(geometry, len(coords)) + return geometry + + +def _validate_positive_int(value: int, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, + (int, np.integer)) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + return int(value) + + +def _validate_modal_data(data: GDataState, operation: str, + dimensions: tuple[int, ...]) -> None: + """Enforce the shared raw-DG input contract for GK projections.""" + if data.num_dims not in dimensions: + expected = " or ".join(f"{dim}-D" for dim in dimensions) + raise ValueError( + f"{operation} requires {expected} data; got {data.num_dims}-D.") + if data.values is None: + raise ValueError(f"{operation} requires a loaded dataset.") + if data.ctx.get("interpolated") or data.ctx.get("value_form", + "modal") != "modal": + raise ValueError(f"{operation} expects un-interpolated modal DG data.") + if not data.ctx.get("basis_type"): + raise ValueError( + f"{operation} requires 'basis_type' metadata on the input data.") + poly_order = data.ctx.get("poly_order") + if isinstance(poly_order, bool) or not isinstance(poly_order, (int, np.integer)) \ + or poly_order < 0: + raise ValueError( + f"{operation} requires a nonnegative integer 'poly_order'.") + if len(data.grid) != data.num_dims or any( + np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 + for axis in data.grid): + raise ValueError( + f"{operation} requires one one-dimensional edge grid per data dimension." + ) + if any(not (np.all(np.diff(axis) > 0) or np.all(np.diff(axis) < 0)) + for axis in data.grid): + raise ValueError( + f"{operation} requires strictly monotonic data edge grids.") + + +def _num_fields(data: GDataState) -> int: + """Return the number of physical fields stored in raw modal data.""" + basis_count = num_basis(data.num_dims, int(data.ctx["poly_order"]), + data.ctx["basis_type"]) + stored = data.values.shape[-1] + if stored % basis_count: + raise ValueError( + f"Data stores {stored} coefficients per cell, which is incompatible " + f"with a {basis_count}-coefficient basis.") + return stored // basis_count + + +def _validate_component(data: GDataState, comp: int) -> int: + if isinstance(comp, bool) or not isinstance(comp, (int, np.integer)): + raise ValueError("comp must be an integer component index.") + comp = int(comp) + num_fields = _num_fields(data) + if not 0 <= comp < num_fields: + raise ValueError( + f"comp {comp} is out of bounds for data with {num_fields} component(s)." + ) + return comp + + +def _interpolation_grid( + data: GDataState) -> tuple[list[np.ndarray], list[np.ndarray]]: + """Return interpolation edges/centers without evaluating field values.""" + num_interp = int(data.ctx["poly_order"]) + 1 + edges = [ + np.linspace(axis[0], axis[-1], + num_interp * (axis.size - 1) + 1) for axis in data.grid + ] + centers = nodal_to_cell_centered_grid( + edges, np.array([axis.size - 1 for axis in edges])) + return edges, centers + + +def _interpolate_component( + data: GDataState, + comp: int) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]: + """Interpolate and return a component already checked by the public API.""" + field = interpolate(data) + cells = field.values.shape[:-1] + centers = nodal_to_cell_centered_grid(field.grid, cells) + return field.grid, centers, field.values[..., comp] + + +def _resample_grid(values: np.ndarray, src_coords: list[np.ndarray], + dst_coords: list[np.ndarray]) -> np.ndarray: + """Linearly resample ``values`` between tensor-product coordinate grids.""" + mesh = np.meshgrid(*dst_coords, indexing="ij") + return RegularGridInterpolator(tuple(src_coords), + values, + bounds_error=False, + fill_value=None)(tuple(mesh)) + + +def _same_grid(left: tuple[np.ndarray, ...] | list[np.ndarray], + right: tuple[np.ndarray, ...] | list[np.ndarray]) -> bool: + return len(left) == len(right) and all( + a.shape == b.shape and np.allclose(a, b, rtol=1e-12, atol=1e-14) + for a, b in zip(left, right)) + + +__all__ = [ + "GKYL_GEOMETRY_ID", + "Geometry", + "geometry_prefix", + "is_geo_mapc2p", + "per_block_path", + "resolve_geometry", +] diff --git a/src/postgkyl/operations/gyrokinetics/rz.py b/src/postgkyl/operations/gyrokinetics/rz.py new file mode 100644 index 00000000..b4b1e051 --- /dev/null +++ b/src/postgkyl/operations/gyrokinetics/rz.py @@ -0,0 +1,359 @@ +"""Project gyrokinetic DG fields onto a physical poloidal R-Z plane.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from scipy.interpolate import PchipInterpolator + +from .geometry import ( + Geometry, + geometry_prefix, + _interpolate_component, + _interpolation_grid, + _resample_grid, + _same_grid, + _validate_component, + _validate_geometry, + _validate_modal_data, + _validate_positive_int, + per_block_path, + resolve_geometry, +) + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +@dataclass(frozen=True) +class RzProjection: + """Precomputed R-Z mapping reusable by fields on one computational grid.""" + + num_dims: int + r: np.ndarray + z: np.ndarray + zc: np.ndarray | None = None + zf: np.ndarray | None = None + box: float | None = None + wind: np.ndarray | None = None + phi0_zf: np.ndarray | None = None + computational_grid: tuple[np.ndarray, ...] | None = None + + +def _fft_poloidal_project(values: np.ndarray, zc: np.ndarray, box: float, + wind: np.ndarray, phi0_zf: np.ndarray, zf: np.ndarray, + phi_tor: float) -> np.ndarray: + """FFT twist-and-shift reconstruction at one physical toroidal angle.""" + nx, ny, nz = values.shape + fk = np.fft.rfft(values, axis=1, norm="forward") + mode_count = fk.shape[1] + + dz = zc[1] - zc[0] + z_extended = np.concatenate(([zc[0] - dz / 2], zc, [zc[-1] + dz / 2])) + fk_extended = np.zeros((nx, mode_count, nz + 2), dtype=complex) + fk_extended[:, :, 1:-1] = fk + phase_shift = (2.0 * np.pi / box) * wind + for mode in range(mode_count): + phase = np.exp(-1j * mode * phase_shift) + fk_extended[:, mode, -1] = 0.5 * (fk[:, mode, -1] + phase * fk[:, mode, 0]) + fk_extended[:, mode, + 0] = 0.5 * (fk[:, mode, 0] + np.conj(phase) * fk[:, mode, -1]) + + fk_zf = (PchipInterpolator(z_extended, fk_extended.real, axis=2)(zf) + + 1j * PchipInterpolator(z_extended, fk_extended.imag, axis=2)(zf)) + + fraction = (phi_tor - phi0_zf) / box + out = np.zeros((nx, len(zf))) + for mode in range(mode_count): + weight = 1.0 if (mode == 0 or + (ny % 2 == 0 and mode == mode_count - 1)) else 2.0 + out += weight * np.real( + fk_zf[:, mode, :] * np.exp(-1j * 2.0 * np.pi * mode * fraction)) + return out + + +def resolve_rz_projection(first: "GDataState", + geo: Geometry, + *, + z_axis: float = 0.0, + nz_interp: int = 8) -> RzProjection: + """Build an R-Z projection for ``first``'s grid and ``geo``. + + Only the computational grid and DG metadata are read from ``first``; + projection construction never evaluates or selects its field values. + ``z_axis`` is the magnetic-axis vertical position in meters. + """ + _validate_modal_data(first, "gk_rz", (2, 3)) + nz_interp = _validate_positive_int(nz_interp, "nz_interp") + _validate_geometry(geo, first.num_dims) + edges, centers = _interpolation_grid(first) + vert_z = geo.vert_z + float(z_axis) + + if first.num_dims == 2: + r = _resample_grid(geo.major_r, geo.coords, edges) + z = _resample_grid(vert_z, geo.coords, edges) + return RzProjection(num_dims=2, + r=r, + z=z, + computational_grid=tuple( + np.array(axis, copy=True) for axis in edges)) + + if geo.phi is None: + raise ValueError( + "The geometry file has no toroidal-angle component; 3-D gk_rz requires one." + ) + + xc, yc, zc = centers + nx, ny, nz = xc.size, yc.size, zc.size + if ny < 2 or nz < 2: + raise ValueError( + "3-D gk_rz requires at least two interpolated y and z points.") + + phi = np.unwrap(np.unwrap(np.unwrap(geo.phi, axis=2), axis=1), axis=0) + phi_field = _resample_grid(phi, geo.coords, centers) + box_estimate = np.mean(np.diff(phi_field[nx // 2, :, nz // 2])) * ny + if not np.isfinite(box_estimate) or np.isclose(box_estimate, 0.0): + raise ValueError( + "Toroidal geometry has a zero or non-finite binormal angular span.") + n0 = max(1, int(round(abs(2.0 * np.pi / box_estimate)))) + box = np.sign(box_estimate) * 2.0 * np.pi / n0 + wind = phi_field[:, 0, -1] - phi_field[:, 0, 0] + + xn, _, zn = edges + zf_edges = np.linspace(zn[0], zn[-1], nz_interp * nz + 1) + zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) + phi0_zf = np.array( + [np.interp(zf, zc, phi_field[ix, 0, :]) for ix in range(nx)]) + + gx, _, gz = geo.coords + r2d = geo.major_r[:, 0, :] + z2d = vert_z[:, 0, :] + gz_rz = gz + if geo.corner is not None: + corner_coords, corner_r, corner_z = geo.corner + if len(corner_coords) != 3: + raise ValueError( + "Corner geometry must be three-dimensional for 3-D gk_rz.") + cx, cz = corner_coords[0], corner_coords[2] + corner_r = corner_r[:, 0, :] + corner_z = corner_z[:, 0, :] + float(z_axis) + r2d = np.concatenate([ + np.interp(gx, cx, corner_r[:, 0])[:, None], r2d, + np.interp(gx, cx, corner_r[:, -1])[:, None] + ], + axis=1) + z2d = np.concatenate([ + np.interp(gx, cx, corner_z[:, 0])[:, None], z2d, + np.interp(gx, cx, corner_z[:, -1])[:, None] + ], + axis=1) + gz_rz = np.concatenate([[cz[0]], gz, [cz[-1]]]) + + r = _resample_grid(r2d, [gx, gz_rz], [xn, zf_edges]) + z = _resample_grid(z2d, [gx, gz_rz], [xn, zf_edges]) + return RzProjection(num_dims=3, + r=r, + z=z, + zc=zc, + zf=zf, + box=box, + wind=wind, + phi0_zf=phi0_zf, + computational_grid=tuple( + np.array(axis, copy=True) for axis in edges)) + + +def _validate_projection(data: "GDataState", projection: RzProjection) -> None: + if projection.num_dims not in (2, 3): + raise ValueError( + f"R-Z projection has invalid dimensionality {projection.num_dims}; expected 2 or 3." + ) + if data.num_dims != projection.num_dims: + raise ValueError( + "Incompatible R-Z projection: projection dimensionality does not match the data." + ) + edges, _ = _interpolation_grid(data) + if projection.computational_grid is not None \ + and not _same_grid(projection.computational_grid, edges): + raise ValueError( + "Incompatible R-Z projection: data computational grid does not match " + "the grid used to build the projection.") + if projection.r.shape != projection.z.shape or projection.r.ndim != 2: + raise ValueError( + "Incompatible R-Z projection: R and Z grids must be matching 2-D arrays." + ) + + if projection.num_dims == 2: + expected = (edges[0].size, edges[1].size) + if projection.r.shape != expected: + raise ValueError( + f"Incompatible R-Z projection: expected grid shape {expected}, " + f"got {projection.r.shape}.") + return + + required = (projection.zc, projection.zf, projection.box, projection.wind, + projection.phi0_zf) + if any(value is None for value in required): + raise ValueError( + "Incompatible R-Z projection: 3-D projection metadata is incomplete.") + if not np.isfinite(projection.box) or np.isclose(projection.box, 0.0): + raise ValueError( + "Incompatible R-Z projection: toroidal angular span must be finite and nonzero." + ) + nx, _, nz = (axis.size - 1 for axis in edges) + if (projection.zc.shape != (nz, ) or projection.wind.shape != (nx, ) + or projection.phi0_zf.shape != (nx, projection.zf.size) + or projection.r.shape != (nx + 1, projection.zf.size + 1)): + raise ValueError( + "Incompatible R-Z projection: projection and data grid shapes differ.") + + +def map_to_rz(data: "GDataState", + projection: RzProjection, + *, + phi_tor: float = 0.0, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Map one component of ``data`` with a reusable ``projection``. + + ``phi_tor`` is in radians and is used only for 3-D field-aligned input. + The source and projection computational grids must be identical. + """ + _validate_modal_data(data, "gk_rz", (2, 3)) + _validate_component(data, comp) + _validate_projection(data, projection) + _, _, values = _interpolate_component(data, comp) + + if projection.num_dims == 2: + out = values[..., np.newaxis] + else: + out = _fft_poloidal_project(values, projection.zc, projection.box, + projection.wind, projection.phi0_zf, + projection.zf, float(phi_tor))[..., np.newaxis] + + return data._result([projection.r, projection.z], + out, + inplace=inplace, + tag=tag, + label=label, + interpolated=True) + + +def rz_projections(datasets, + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + z_axis: float = 0.0, + nz_interp: int = 8) -> dict[str | None, RzProjection]: + """Build one reusable R-Z projection per block geometry. + + Frames from the same block share a projection; distinct blocks resolve + their own geometry. A ``'*'`` in an explicit geometry path is replaced by + the dataset's block index. + """ + projections: dict[str | None, RzProjection] = {} + for data in datasets: + key = geometry_prefix(data.file_name) + if key in projections: + continue + block = data.ctx.get("block") + geometry = resolve_geometry(data.file_name, + mapc2p=per_block_path(mapc2p, block), + nodes_file=per_block_path(nodes_file, block)) + projections[key] = resolve_rz_projection(data, + geometry, + z_axis=z_axis, + nz_interp=nz_interp) + return projections + + +def projection_for(projections: dict[str | None, RzProjection], + data: "GDataState") -> RzProjection: + """Return the projection belonging to ``data``'s block.""" + return projections[geometry_prefix(data.file_name)] + + +def gk_rz( + data: "GDataState", + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + z_axis: float = 0.0, + phi_tor: float = 0.0, + nz_interp: int = 8, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None, +) -> "GDataState": + """Interpolate one DG component and project it onto a physical R-Z grid. + + ``data`` must be un-interpolated modal DG data with two computational + dimensions, or three for a field-aligned reconstruction. Geometry is + inferred from ``data.file_name``: the pointwise + ``'-geo_int_nodes.gkyl'`` file is preferred, falling back to + ``'-geo_int_mapc2p.gkyl'``. ``nodes_file`` or ``mapc2p`` may + override that choice, but they are mutually exclusive; ``mapc2p=''`` + forces the inferred modal filename. + + Args: + data: Un-interpolated 2-D or 3-D modal DG field. + mapc2p: Optional explicit modal geometry path, or ``''`` for inferred. + nodes_file: Optional explicit nodal geometry path. + z_axis: Magnetic-axis vertical position in meters, added to geometry Z. + phi_tor: Toroidal angle in radians for a 3-D poloidal reconstruction. + nz_interp: Positive integer z-direction up-sampling factor for 3-D data. + comp: Zero-based physical field component to map (default first). + inplace: Replace ``data`` rather than returning a new concrete instance. + tag: Optional result tag; ``None`` preserves the source tag. + label: Optional result label; ``None`` preserves the source label. + + Returns: + The caller's concrete data class, marked ``interpolated=True``. If the + interpolated input counts are ``(Nx, Nz)``, 2-D grid arrays have shape + ``(Nx+1, Nz+1)`` and values ``(Nx, Nz, 1)``. For interpolated 3-D counts + ``(Nx, Ny, Nz)``, grid arrays have shape + ``(Nx+1, nz_interp*Nz+1)`` and values + ``(Nx, nz_interp*Nz, 1)``. + + Raises: + ValueError: For mutually exclusive or missing geometry, input other than + un-interpolated 2-D/3-D modal data, a missing 3-D toroidal angle, + invalid ``comp`` or ``nz_interp``, malformed geometry, or a projection + incompatible with its data grid (when using :func:`map_to_rz`). + """ + _validate_modal_data(data, "gk_rz", (2, 3)) + _validate_positive_int(nz_interp, "nz_interp") + _validate_component(data, comp) + geometry = resolve_geometry(data.file_name, + mapc2p=mapc2p, + nodes_file=nodes_file) + projection = resolve_rz_projection(data, + geometry, + z_axis=z_axis, + nz_interp=nz_interp) + return map_to_rz(data, + projection, + phi_tor=phi_tor, + comp=comp, + inplace=inplace, + tag=tag, + label=label) + + +__all__ = [ + "Geometry", + "RzProjection", + "geometry_prefix", + "gk_rz", + "map_to_rz", + "per_block_path", + "projection_for", + "resolve_geometry", + "resolve_rz_projection", + "rz_projections", +] diff --git a/src/postgkyl/operations/info.py b/src/postgkyl/operations/info.py new file mode 100644 index 00000000..298765aa --- /dev/null +++ b/src/postgkyl/operations/info.py @@ -0,0 +1,20 @@ +"""The ``info`` verb -- print/return summaries for one or more datasets.""" + +from __future__ import annotations + +from postgkyl.gdatastate import flatten_datasets +from postgkyl.gdatastate.gdatastate import GDataState + + +def info(*datasets: GDataState, no_header: bool = False) -> list: + """Print a summary for each dataset; return the list of summary strings. + + Accepts ``info(a, b)`` or ``info([a, b])``. Each dataset's own ``info`` method + (a pure state reader on the container) does the formatting. + + Args: + datasets: Datasets whose summaries are returned. + no_header: Omit the descriptive heading from every summary. + """ + states = flatten_datasets(datasets) + return [d.info(index=i, no_header=no_header) for i, d in enumerate(states)] diff --git a/src/postgkyl/operations/integrate.py b/src/postgkyl/operations/integrate.py new file mode 100644 index 00000000..b102106b --- /dev/null +++ b/src/postgkyl/operations/integrate.py @@ -0,0 +1,244 @@ +"""Full and partial grid integration under one ``integrate`` verb. + +The return shape states which operation was requested: + +* integrating every spatial direction is terminal and returns one number per + field; +* integrating a strict subset returns a dataset over the surviving + directions. + +Modal data never leave Gkeyll. Full integration uses +``gkyl_array_integrate`` directly. Partial integration uses +``gkyl_array_average`` and scales its modal result by the physical volume of +the removed directions, which is exactly ``int f dx^axes`` rather than a +sampled/trapezoidal approximation. Point-value data use the NumPy integration +path at their true point locations. + +A curvilinear axis -- part of a joint, non-separable ``.map(space="conf")`` +block -- has no meaningful independent width. Point-value integration must +therefore remove the whole mapped block at once, using its physical cell +volumes (the Jacobian-determinant change-of-variables weight). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl import dg +from postgkyl.gdatastate import materialize_point_values +from postgkyl.numerics import calculus, curvilinear + +from ._curvilinear import curvilinear_blocks + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _parse_axes(axis: int | tuple | str | None, ndim: int) -> tuple[int, ...]: + axes = tuple(int(a) for a in calculus.parse_axis(axis, ndim)) + if not axes: + raise ValueError("integrate needs at least one axis") + if len(set(axes)) != len(axes): + raise ValueError(f"integrate axes must be distinct, got {axes}") + if min(axes) < 0 or max(axes) >= ndim: + raise ValueError(f"integrate axes {axes} out of range for a {ndim}D field") + return tuple(sorted(axes)) + + +def _native_basis(data: "GDataState") -> tuple[str, int]: + if data.backend != "gkyl": + raise ValueError( + "exact DG integration needs native modal data and is not available " + "without the Gkeyll library") + if data.ctx.get("value_form", "modal") != "modal": + raise ValueError(f"exact DG integration expects the modal value_form, not " + f"'{data.ctx['value_form']}'; call .to_modal() first") + basis_type = data.ctx.get("basis_type") + poly_order = data.ctx.get("poly_order") + if basis_type is None or poly_order is None: + raise ValueError("dataset has no basis_type/poly_order metadata") + return str(basis_type), int(poly_order) + + +def _native_grid(data: "GDataState") -> dict: + return { + "ndim": data.num_dims, + "lower": np.asarray(data.ctx["lower"]), + "upper": np.asarray(data.ctx["upper"]), + "cells": np.asarray(data.ctx["cells"]), + } + + +def _native_full(data: "GDataState", op: str): + basis_type, poly_order = _native_basis(data) + result = dg.modal.integrate(_native_grid(data), + basis_type, + poly_order, + data.native, + op=op) + return float(result[0]) if result.size == 1 else result + + +def _native_partial(data: "GDataState", axes: tuple[int, ...], *, inplace: bool, + tag: str | None, label: str | None): + basis_type, poly_order = _native_basis(data) + grid = _native_grid(data) + keep_dirs, cells, out = dg.modal.average(grid, basis_type, data.num_dims, + poly_order, data.native, axes) + + # gkyl_array_average returns int(f dx^axes) / int(dx^axes). Scaling the + # reduced modal coefficients recovers the integral exactly and stays native. + lengths = grid["upper"] - grid["lower"] + out = dg.modal.scale(out, float(np.prod(lengths[list(axes)]))) + new_grid = [np.asarray(data.grid[d]) for d in keep_dirs] + return data._result(new_grid, + out, + inplace=inplace, + tag=tag, + label=label, + cells=np.asarray(cells)) + + +def _point_integral(data: "GDataState", axes: tuple[int, ...]): + data._require_operable() + shadow = materialize_point_values(data) + grid = list(shadow.grid) + values = shadow.values + + blocks = curvilinear_blocks(grid, data.ctx.get("mapped_axes", {})) + requested = set(axes) + curvilinear_runs = [] + handled = set() + for off, dims in blocks.items(): + overlap = requested & set(dims) + if not overlap: + continue + if overlap != set(dims): + raise ValueError( + f"integrate: axis/axes {sorted(overlap)} belong to a curvilinear " + f"(mapped) block spanning dimensions {dims}; a partial reduction " + "of the block has no single physical answer -- include every " + "axis of the block together in the same call") + curvilinear_runs.append((off, dims)) + handled.update(dims) + + separable_axes = tuple(a for a in axes if a not in handled) + if separable_axes: + grid, values = calculus.integrate(grid, values, separable_axes) + + for _, dims in curvilinear_runs: + m = len(dims) + block_coords = [grid[d] for d in dims] + volume = curvilinear.cell_volume(block_coords) + volume = volume.reshape(volume.shape + (1, ) * (values.ndim - m)) + moved = np.moveaxis(values, dims, range(m)) + reduced = np.sum(moved * volume, axis=tuple(range(m)), keepdims=True) + values = np.moveaxis(reduced, range(m), dims) + for d in dims: + grid[d] = np.array([grid[d].mean()]) + return grid, values + + +def _terminal_value(values: np.ndarray): + result = np.asarray(values).reshape(-1, values.shape[-1])[0] + return float(result[0]) if result.size == 1 else np.array(result, copy=True) + + +def _remaining_mapped_axes(data: "GDataState", keep_dirs: list[int]) -> dict: + old = data.ctx.get("mapped_axes", {}) + old_to_new = {old_dim: new_dim for new_dim, old_dim in enumerate(keep_dirs)} + groups: dict[int, list[int]] = {} + for old_dim, offset in old.items(): + if old_dim in old_to_new: + groups.setdefault(offset, []).append(old_dim) + + result = {} + for old_dims in groups.values(): + new_dims = [old_to_new[d] for d in old_dims] + new_offset = min(new_dims) + result.update({d: new_offset for d in new_dims}) + return result + + +def _require_partial_options(*, inplace: bool, tag: str | None, + label: str | None) -> None: + if inplace or tag is not None or label is not None: + raise ValueError( + "inplace, tag, and label apply only to partial integration, which " + "returns a dataset") + + +def integrate(data: "GDataState", + axis: int | tuple | str | None = None, + *, + op: str = "none", + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Integrate over all or a subset of a dataset's spatial axes. + + Integrating every axis is terminal and returns one number per field. + Integrating a strict subset returns a dataset over the surviving axes. + Native modal input uses Gkeyll for both cases and remains modal/native after + a partial integration; no interpolation or point sampling occurs. Nodal, + quadrature, and NumPy-backed inputs are integrated numerically at their true + point locations. + + Args: + data: Dataset to integrate. Modal data are integrated exactly in DG space; + point-value data use their physical point grid. + axis: Axis or axes to integrate: an integer, tuple of integers, + comma-separated string (``"0,1"``), colon slice string (``"0:2"``), or + ``None`` for every spatial axis. + op: Full native-DG integrand operation: ``"none"``, ``"abs"``, or + ``"sq"``. Partial and point-value integration support ``"none"`` only. + inplace: Mutate ``data`` for partial integration; unavailable for a full, + terminal integration. + tag: Optional tag for a partial-integration result. + label: Optional label for a partial-integration result. + + Returns: + A float (one field) or NumPy array (multiple fields) when every axis is + integrated; otherwise a dataset over the surviving axes. A partial modal + result is exact, modal, and Gkeyll-native. + + Raises: + ValueError: If the axes are invalid; a partial curvilinear mapped block is + requested; ``op`` is used outside a full native-DG integration; or + dataset-only result options are used for a terminal integration. + """ + axes = _parse_axes(axis, data.num_dims) + full = len(axes) == data.num_dims + modal = (data.backend == "gkyl" + and data.ctx.get("value_form", "modal") == "modal") + + if full: + _require_partial_options(inplace=inplace, tag=tag, label=label) + if modal: + return _native_full(data, op) + if op != "none": + raise ValueError("op is available only for full native-DG integration") + _, values = _point_integral(data, axes) + return _terminal_value(values) + + if op != "none": + raise ValueError("op is available only for full native-DG integration") + if modal: + return _native_partial(data, axes, inplace=inplace, tag=tag, label=label) + + grid, values = _point_integral(data, axes) + keep_dirs = [d for d in range(data.num_dims) if d not in axes] + values = np.squeeze(values, axis=axes) + new_grid = [grid[d] for d in keep_dirs] + mapped_axes = _remaining_mapped_axes(data, keep_dirs) + return data._result(new_grid, + values, + inplace=inplace, + tag=tag, + label=label, + interpolated=True, + value_form=None, + mapped_axes=mapped_axes, + grid_type="mapped" if mapped_axes else "uniform") diff --git a/src/postgkyl/operations/interpolate.py b/src/postgkyl/operations/interpolate.py new file mode 100644 index 00000000..134d2da7 --- /dev/null +++ b/src/postgkyl/operations/interpolate.py @@ -0,0 +1,62 @@ +"""The ``interpolate`` verb -- DG coefficients -> values on a uniform mesh.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import dg + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def interpolate(data: "GDataState", + *, + num_interp: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Interpolate DG (modal/nodal) data onto a uniform evaluation mesh. + + Basis, polynomial order, and value_form are properties of ``data`` itself, + fixed at load time (``pg.load(..., basis_type=..., poly_order=..., + value_form=...)`` or the CLI's ``-b``/``-p``/``-v``) -- this verb only + ever reads them off ``data.ctx``. The result is flagged + ``interpolated=True`` so it becomes safe for element-wise math. + + Args: + data: Dataset containing DG coefficients. + num_interp: Evaluation points per cell; use the basis default when omitted. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + basis_type = data.ctx.get("basis_type") + if not basis_type: + raise ValueError( + "dataset has no 'basis_type' metadata; set it at load time " + "(pg.load(..., basis_type=...) or the CLI's -b/--basis).") + + poly_order = data.ctx.get("poly_order") + if poly_order is None: + raise ValueError( + "dataset has no 'poly_order' metadata; set it at load time " + "(pg.load(..., poly_order=...) or the CLI's -p/--poly_order).") + + value_form = data.ctx.get("value_form", "modal") + if data.backend == "gkyl" and value_form != "modal": + raise ValueError(f"interpolate expects the modal value_form, not " + f"'{value_form}'; call .to_modal() first.") + + grid, values = dg.interpolate(data.values, + data.grid, + poly_order=poly_order, + basis_type=basis_type, + nodal=(value_form == "nodal"), + num_interp=num_interp) + return data._result(grid, + values, + inplace=inplace, + tag=tag, + label=label, + interpolated=True) diff --git a/src/postgkyl/operations/local_poly.py b/src/postgkyl/operations/local_poly.py new file mode 100644 index 00000000..b991454a --- /dev/null +++ b/src/postgkyl/operations/local_poly.py @@ -0,0 +1,67 @@ +"""The ``local_poly`` verb -- modal DG coefficients -> a discontinuity- +preserving plotting mesh (see ``dg.local_poly``).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import dg + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def local_poly(data: "GDataState", + *, + npoints: int = 2, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Evaluate the DG polynomial cell-by-cell onto a plotting mesh that keeps + every inter-cell discontinuity visible, instead of the continuous refined + mesh ``interpolate`` produces. + + ``npoints`` reference points span the whole cell (``[-1, 1]``, endpoints + included) and a NaN is spliced in at every cell interface, so a plot breaks + the curve there rather than drawing a spuriously smooth line across it. + + Basis, polynomial order, and value_form are properties of ``data`` itself, + fixed at load time, same as ``interpolate``. The result is flagged + ``interpolated=True`` so it becomes safe for element-wise math. + + Args: + data: Dataset containing modal DG coefficients. + npoints: Evaluation points in each cell and direction. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + basis_type = data.ctx.get("basis_type") + if not basis_type: + raise ValueError( + "dataset has no 'basis_type' metadata; set it at load time " + "(pg.load(..., basis_type=...) or the CLI's -b/--basis).") + + poly_order = data.ctx.get("poly_order") + if poly_order is None: + raise ValueError( + "dataset has no 'poly_order' metadata; set it at load time " + "(pg.load(..., poly_order=...) or the CLI's -p/--poly_order).") + + value_form = data.ctx.get("value_form", "modal") + if data.backend == "gkyl" and value_form != "modal": + raise ValueError(f"local_poly expects the modal value_form, not " + f"'{value_form}'; call .to_modal() first.") + + grid, values = dg.local_poly(data.values, + data.grid, + poly_order=poly_order, + basis_type=basis_type, + nodal=(value_form == "nodal"), + npoints=npoints) + return data._result(grid, + values, + inplace=inplace, + tag=tag, + label=label, + interpolated=True) diff --git a/src/postgkyl/operations/magsq.py b/src/postgkyl/operations/magsq.py new file mode 100644 index 00000000..61f28c7f --- /dev/null +++ b/src/postgkyl/operations/magsq.py @@ -0,0 +1,44 @@ +"""The ``magsq`` verb -- magnitude squared of a vector field.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import numerics + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def magsq(data: "GDataState", + *, + coords: str = "0:3", + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Magnitude squared of a vector field. + + Sums the squares of the selected components (``numerics.mag_sq``), + returning a single-component field. + + Args: + data: the dataset holding the vector field; must be NumPy-backed. + coords: ``"start:end"`` slice of the component axis to sum the squares + of. Defaults to the first three components. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A single-component dataset of the magnitude squared. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed). + """ + if data.backend == "gkyl": + raise ValueError( + "magsq operates on interpolated (NumPy) values; call .interpolate() " + "first -- summing squares of raw DG coefficients would mix basis functions." + ) + grid, values = numerics.mag_sq(data.grid, data.values, coords=coords) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/operations/map.py b/src/postgkyl/operations/map.py new file mode 100644 index 00000000..17244948 --- /dev/null +++ b/src/postgkyl/operations/map.py @@ -0,0 +1,163 @@ +"""The ``map`` verb -- deform a dataset's grid by evaluating a coordinate map. + +See ``MAPPING.md`` for the full design. A mapping file is a DG field whose +components hold the coefficients of the physical coordinates of each mapped +dimension; this verb evaluates those coefficients at the *target*'s own grid +points (:func:`postgkyl.dg.map_grid`) and splices the resulting arrays into +a copy of the target's grid. Only the grid changes -- the mapping's +coefficients are read straight from its native modal storage and are never +interpolated, and the target's values are passed through unchanged (no +copy: this verb never touches them). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import dg +from postgkyl.gdatastate.gdatastate import GDataState + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState as _GDataState + + +def map(data: "_GDataState", + mapping: "str | _GDataState", + *, + space: str = "conf", + basis_type: str | None = None, + poly_order: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "_GDataState": + """Replace a block of ``data``'s grid axes with mapped coordinates. + + Evaluates the mapping's DG coefficients at ``data``'s existing grid + points (no resolution parameter, no alignment arithmetic -- the mapped + axes always keep the shape of the axes they replace) and splices the + result into a copy of ``data``'s grid. + + Args: + data: The dataset whose grid is deformed; must be NumPy-backed + (post-``interpolate()``), like ``select``. + mapping: The coordinate-mapping field, as a filename or an + already-loaded dataset. Read from its native modal coefficients -- + never interpolated. Its number of dimensions (``m``) sets how many + of ``data``'s axes are replaced. + space: ``'conf'`` deforms the leading ``m`` axes (offset 0), + curvilinearly (every physical coordinate is evaluated over all ``m`` + mapped dimensions, so non-separable maps such as rotations work); its + component count must be ``m * num_basis`` for an ``m``-D basis. + ``'vel'`` deforms the trailing ``m`` axes (offset + ``data.num_dims - m``); Gkeyll's velocity-space maps + (``mapc2p_vel``) are diagonal -- each dimension is evaluated by its + *own* 1-D map, so the component count must be ``m * num_basis`` for a + *1-D* basis. For a combined map, apply the verb twice. + basis_type: ``mapping``'s ``basis_type`` (long name, e.g. + ``"serendipity"``), set at the moment this call loads it -- only + takes effect when ``mapping`` is given as a filename (a property of + already-loaded data can't be re-specified here). Velocity-space + mapping files (``mapc2p_vel``) commonly carry no basis metadata at + all, so this is typically required for ``space="vel"``. + poly_order: ``mapping``'s ``poly_order``, set the same way as + ``basis_type``. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset carrying the deformed grid; ``ctx["grid_type"]`` is set to + ``"mapped"`` and ``ctx["mapped_axes"]`` records, for every absolute + dimension touched so far (by this call and any earlier one), the + ``offset`` of the mapped block it belongs to -- ``select``'s + curvilinear guard needs this to convert an absolute dimension index + back to the curvilinear grid array's own (relative) axis. The values + array is untouched. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed); if ``space`` is + neither ``'conf'`` nor ``'vel'``; if the map does not fit ``data``'s + dimensionality; if the mapping has no ``basis_type``/``poly_order`` + metadata (and none was given at load time); or if its component + count does not match the expected ``m * num_basis``. + """ + if data.backend == "gkyl": + raise ValueError( + "map operates on interpolated (NumPy) target grids; call .interpolate() " + "first -- deforming a native modal grid has no basis-space meaning.") + + # basis_type/poly_order are load-time properties of the mapping dataset; + # they only apply here when this call is the one loading it (a filename), + # threaded straight into GDataState's own override mechanism -- the one + # home for "correct a file's missing/mislabeled basis metadata." + map_data = (mapping if isinstance(mapping, GDataState) else GDataState( + mapping, basis_type=basis_type, poly_order=poly_order)) + m = map_data.num_dims + num_dims = data.num_dims + + if space == "conf": + offset = 0 + elif space == "vel": + offset = num_dims - m + else: + raise ValueError(f"map: 'space' must be 'conf' or 'vel', got {space!r}.") + + if offset < 0 or offset + m > num_dims: + raise ValueError( + f"map: a {m}D {space} map does not fit a {num_dims}D dataset.") + + resolved_basis_type = map_data.ctx.get("basis_type") + resolved_poly_order = map_data.ctx.get("poly_order") + if resolved_basis_type is None or resolved_poly_order is None: + raise ValueError( + "map: the mapping dataset has no 'basis_type'/'poly_order' " + "metadata; pass basis_type=.../poly_order=... (mapping as a " + "filename), or load it explicitly first with those set.") + + # Velocity-space maps (mapc2p_vel) are diagonal: each mapped dimension is + # its own separate 1-D map, so its basis is 1-D regardless of m; a + # configuration-space map (mapc2p/mc2nu) is one joint m-D curvilinear map. + basis_dim = 1 if space == "vel" else m + num_basis = dg.num_basis(basis_dim, resolved_poly_order, resolved_basis_type) + if map_data.num_comps != m * num_basis: + raise ValueError( + f"map: mapping has {map_data.num_comps} component(s), expected " + f"m * num_basis = {m} * {num_basis} = {m * num_basis} for a " + f"{basis_dim}D {resolved_basis_type} p{resolved_poly_order} map" + + (" per velocity dimension." if space == "vel" else ".")) + + target_axes = list(data.grid[offset:offset + m]) + map_ctx = { + "lower": map_data.ctx["lower"], + "upper": map_data.ctx["upper"], + "cells": map_data.ctx["cells"], + "basis_type": resolved_basis_type, + "poly_order": resolved_poly_order, + "value_form": map_data.ctx.get("value_form", "modal"), + } + if space == "vel": + new_axes = dg.map_grid_separable(map_data.get_values(), map_ctx, + target_axes) + else: + new_axes = dg.map_grid(map_data.get_values(), map_ctx, target_axes) + + grid = list(data.grid) + for d in range(m): + grid[offset + d] = new_axes[d] + + # Record, per absolute dimension, the offset of the mapped block it + # belongs to -- a curvilinear (m > 1) grid array's own axis k corresponds + # to absolute dimension offset + k, not to the array's position in + # `grid`, so `select`'s curvilinear guard needs this to convert back. + # Merge with any prior block (e.g. a separate `space="vel"` map applied + # after a `space="conf"` one) rather than overwrite it. + mapped_axes = dict(data.ctx.get("mapped_axes", {})) + mapped_axes.update({offset + d: offset for d in range(m)}) + + return data._result(grid, + data.values, + inplace=inplace, + tag=tag, + label=label, + grid_type="mapped", + mapped_axes=mapped_axes) diff --git a/src/postgkyl/operations/mask.py b/src/postgkyl/operations/mask.py new file mode 100644 index 00000000..ce294797 --- /dev/null +++ b/src/postgkyl/operations/mask.py @@ -0,0 +1,82 @@ +"""The ``mask`` verb -- mask out values by a mask dataset or by thresholds.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def mask(data: "GDataState", + mask_data: "GDataState | None" = None, + *, + lower: float | None = None, + upper: float | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Mask out values using a mask dataset or numeric thresholds. + + Returns a dataset whose values are a ``numpy.ma`` masked array. Exactly + one of the masking modes is applied, with ``mask_data`` taking precedence: + + - ``mask_data``: mask cells where the mask dataset's field is negative, + repeated across ``data``'s components. Load the mask field yourself + (e.g. ``pg.load(mask_path)``) -- this verb takes data, never a file path + (``operations`` never touches ``io``). + - ``lower`` and ``upper``: mask values outside the closed range + ``[lower, upper]``. + - ``lower`` only: mask values below ``lower``. + - ``upper`` only: mask values above ``upper``. + + Args: + data: the dataset to mask; must be NumPy-backed. + mask_data: an already-loaded dataset whose field selects the mask + (negative -> masked); it must have exactly one component -- the mask + is broadcast across every component of ``data`` via + ``np.repeat(mask_data.values, data.num_comps, axis=-1)``, which only + produces a shape matching ``data.values`` when ``mask_data`` is + single-component. A multi-component ``mask_data`` raises from the + subsequent ``np.ma.masked_where`` broadcast, not from an explicit + check here. + lower: lower threshold. Combined with ``upper`` masks outside the range; + alone masks values below it. + upper: upper threshold. Combined with ``lower`` masks outside the range; + alone masks values above it. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset whose values are a masked array. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or if none of + ``mask_data``, ``lower``, or ``upper`` is provided. + IndexError: if ``mask_data`` has more than one component -- the + repeated mask no longer matches ``data.values``'s shape and + ``np.ma.masked_where`` rejects the mismatched condition array. + """ + if data.backend == "gkyl": + raise ValueError( + "mask operates on interpolated (NumPy) values; call .interpolate() " + "first -- masking raw DG coefficients has no basis-space meaning.") + values = data.values + if mask_data is not None: + mask_field = mask_data.values + mask_rep = np.repeat(mask_field, data.num_comps, axis=-1) + masked = np.ma.masked_where(mask_rep < 0.0, values) + elif lower is not None and upper is not None: + masked = np.ma.masked_outside(values, lower, upper) + elif lower is not None: + masked = np.ma.masked_less(values, lower) + elif upper is not None: + masked = np.ma.masked_greater(values, upper) + else: + raise ValueError( + "mask: no masking information specified (provide mask_data, lower, " + "or upper).") + return data._result(data.grid, masked, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/operations/print.py b/src/postgkyl/operations/print.py new file mode 100644 index 00000000..5304ae01 --- /dev/null +++ b/src/postgkyl/operations/print.py @@ -0,0 +1,32 @@ +"""Print stored values or grid coordinates for one or more datasets.""" + +from __future__ import annotations + +import builtins + +import numpy as np + +from postgkyl.gdatastate import flatten_datasets +from postgkyl.gdatastate.gdatastate import GDataState + + +def print(*datasets: GDataState, + use: str | None = None, + grid: bool = False) -> None: + """Print the values (or grid) of the selected datasets. + + Values are printed as stored, with singleton dimensions squeezed and + 16-digit precision. For modal data these are DG coefficients; interpolate + first to print field values. Printing leaves the datasets unchanged. + + Args: + datasets: Datasets, groups, or nested iterables of datasets to print. + use: Select only datasets carrying this tag. + grid: Print each grid axis instead of the values. + """ + for data in flatten_datasets(datasets): + if use is not None and data.tag != use: + continue + arrays = data.grid if grid else (np.asarray(data.values).squeeze(), ) + for array in arrays: + builtins.print(np.array2string(array, precision=16)) diff --git a/src/postgkyl/operations/relchange.py b/src/postgkyl/operations/relchange.py new file mode 100644 index 00000000..fcd429ff --- /dev/null +++ b/src/postgkyl/operations/relchange.py @@ -0,0 +1,54 @@ +"""The ``relchange`` verb -- relative change between two datasets.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import numerics + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _require_field_domain(data: "GDataState", who: str) -> None: + if data.backend == "gkyl": + raise ValueError( + f"relchange operates on interpolated (NumPy) values; call .interpolate() " + f"first on {who} -- dividing raw DG coefficients would mix basis functions." + ) + + +def relchange(data0: "GDataState", + data: "GDataState", + *, + comp: int | str | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Relative change of ``data`` with respect to the baseline ``data0``. + + Computes ``(data - data0) / data0`` component-wise (``numerics.rel_change``). + Both datasets are assumed to share the same grid and component layout. + + Args: + data0: the baseline ("before") dataset -- the denominator. + data: the dataset whose relative change is computed; the returned + dataset is built from this one (its grid/ctx are the base of the + result). + comp: when given, every numerator component is divided by this single + baseline component instead of its own (e.g. normalize every energy + component by the total energy component). None divides component-wise. + inplace: mutate and return ``data`` instead of a new dataset. + tag: optional tag for the returned dataset. + label: optional label for the returned dataset. + + Returns: + A dataset of the relative change, built from ``data``. + + Raises: + ValueError: if either operand is native modal (gkyl-backed). + """ + _require_field_domain(data0, "'data0'") + _require_field_domain(data, "'data'") + grid, values = numerics.rel_change(data.grid, data0.values, data.values, comp) + return data._result(grid, values, inplace=inplace, tag=tag, label=label) diff --git a/src/postgkyl/operations/represent.py b/src/postgkyl/operations/represent.py new file mode 100644 index 00000000..710344ea --- /dev/null +++ b/src/postgkyl/operations/represent.py @@ -0,0 +1,132 @@ +"""The value_form verbs -- explicit modal · nodal · quad changes + ``apply``. + +Conversions are **never implicit** (REFACTOR_GKEYLL_FFI.md §3b): these verbs are +the only way a dataset changes value_form, and each one stamps +``ctx["value_form"]`` (and ``ctx["num_quad"]`` for quad data) so ``info`` +always shows what the numbers mean. All of them keep the data gkyl-native. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from postgkyl import dg + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + +VALUE_FORMS = ("modal", "nodal", "quad") + + +def _native_basis(data: "GDataState"): + """(basis_type, ndim, poly_order) for a gkyl-backed dataset, or raise.""" + if data.backend != "gkyl": + raise ValueError( + "value_form changes act on native (gkyl-backed) DG data; " + "this dataset is NumPy-backed (already interpolated, or loaded " + "without the Gkeyll library).") + basis_type = data.ctx.get("basis_type") + poly_order = data.ctx.get("poly_order") + if basis_type is None or poly_order is None: + raise ValueError("dataset has no basis_type/poly_order metadata") + return str(basis_type), data.num_dims, int(poly_order) + + +def represent(data: "GDataState", + *, + to: str, + num_quad: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Convert a native dataset to the ``to`` value_form (explicitly). + + ``modal`` <-> ``nodal`` is exact; ``modal`` -> ``quad`` evaluates at + ``num_quad`` (default ``p+1``) Gauss–Legendre points per dimension; + ``quad`` -> ``modal`` projects back with the rule the data was made with. + ``nodal`` <-> ``quad`` composes through modal. + + Args: + data: Native dataset to convert. + to: Target value representation. + num_quad: Gauss points per direction for quadrature representation. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + if to not in VALUE_FORMS: + raise ValueError(f"unknown value_form '{to}'; " + f"choices: {VALUE_FORMS}") + basis_type, ndim, poly_order = _native_basis(data) + cur = data.ctx.get("value_form", "modal") + arr = data.native + + if cur != to: + if cur == "nodal": # leave nodal (exact) + arr = dg.rep.nodal_to_modal(basis_type, ndim, poly_order, arr) + elif cur == "quad": # leave quad (projection, with the data's own rule) + nq = data.ctx.get("num_quad") + if nq is None: + raise ValueError("quad-represented dataset lost its 'num_quad' ctx") + arr = dg.rep.quad_to_modal(basis_type, ndim, poly_order, arr, int(nq)) + # arr is now modal + if to == "nodal": + arr = dg.rep.modal_to_nodal(basis_type, ndim, poly_order, arr) + elif to == "quad": + nq = int(num_quad) if num_quad else poly_order + 1 + arr = dg.rep.modal_to_quad(basis_type, ndim, poly_order, arr, nq) + else: + arr = arr.clone() + + return data._result(data.grid, + arr, + inplace=inplace, + tag=tag, + label=label, + value_form=to, + num_quad=(int(num_quad) if num_quad else poly_order + + 1) if to == "quad" else None) + + +def apply(data: "GDataState", + fn, + *, + num_quad: int | None = None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Apply ``fn`` pointwise via quadrature: modal -> quad -> fn -> modal. + + The explicit spelling of nonlinear pointwise operations on DG data (e.g. + ``d.apply(np.sqrt)``): evaluate at ``num_quad`` (default ``p+1``) Gauss + points, apply ``fn`` to the values, project back onto the basis. The result + stays modal and gkyl-native; the projection is exact when ``fn(f)·b_j`` has + degree ≤ 2·num_quad−1 -- raise ``num_quad`` to de-alias. + + Args: + data: Native modal dataset to transform. + fn: Python callable applied to the quadrature-point values. + num_quad: Gauss points per dimension; defaults to ``poly_order + 1``. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Returns: + A modal dataset containing the projected pointwise result. + + Raises: + ValueError: If ``data`` is not native modal data or lacks basis metadata. + """ + basis_type, ndim, poly_order = _native_basis(data) + if data.ctx.get("value_form", "modal") != "modal": + raise ValueError("apply() expects modal data; call .to_modal() first.") + nq = int(num_quad) if num_quad else poly_order + 1 + out = dg.rep.apply_pointwise(basis_type, ndim, poly_order, data.native, fn, + nq) + return data._result(data.grid, + out, + inplace=inplace, + tag=tag, + label=label, + applied=getattr(fn, "__name__", str(fn)), + applied_num_quad=nq) diff --git a/src/postgkyl/operations/select.py b/src/postgkyl/operations/select.py new file mode 100644 index 00000000..e7f1d628 --- /dev/null +++ b/src/postgkyl/operations/select.py @@ -0,0 +1,203 @@ +"""The ``select`` verb -- subselect coordinates and components.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl import dg +from postgkyl.numerics import idx_parser + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _curvilinear_coord_curve(grid_arr: np.ndarray, rel: int, d: int, + offset: int, values_shape: tuple, + touched: set) -> np.ndarray: + """A 1-D coordinate curve along ``grid_arr``'s relative axis ``rel``. + + Holds every other axis fixed at cell 0. That is exact only if + ``grid_arr`` doesn't actually vary along those other axes -- true for the + diagonal/separable maps this codebase's real mapped grids tend to produce + (e.g. field-aligned ``mc2nu`` coordinates, each depending on only one + computational axis even though stored jointly), false for a genuinely + coupled map (e.g. a rotation). Raises rather than silently picking an + arbitrary cross-section when a still-*unresolved* other axis of the + block demonstrably varies the coordinate. An axis is resolved -- and + skipped by the check -- either because ``values_shape`` (the dataset's + own values array; the one true record of what a prior, possibly separate, + ``select()`` call already narrowed) already holds a single cell on its + absolute dimension ``offset + k``, or because a selector targeted it + earlier in this same call (``touched``, before ``values`` itself was + re-sliced at the end of that call). Either way the caller has made a + deliberate choice of cross-section that this search takes as given. + """ + fixed = tuple(0 if k != rel else slice(None) for k in range(grid_arr.ndim)) + full_range = np.ptp(grid_arr) + for k in range(grid_arr.ndim): + if k == rel or k in touched or values_shape[offset + k] == 1: + continue + if np.ptp(grid_arr, axis=k).max() > 1e-9 * full_range: + raise ValueError( + f"select: z{d}'s physical coordinate also varies along another " + "axis of the same mapped (curvilinear) block, so a coordinate " + "value or slice string has no single answer -- select an " + "integer index for that other axis first (narrowing it to one " + f"cell), or pass an integer index for z{d} itself.") + return grid_arr[fixed] + + +def select(data: "GDataState", + *, + comp=None, + z0=None, + z1=None, + z2=None, + z3=None, + z4=None, + z5=None, + inplace: bool = False, + tag: str | None = None, + label: str | None = None): + """Select part of a dataset by coordinate (``z0``-``z5``) and/or component. + + Each selector accepts an int index, a float coordinate value, or a slice + string ``"start:end"``; ``comp`` additionally accepts ``"a,b"``. Unspecified + axes are kept in full. The selected dimension is retained (length-1), matching + the legacy behaviour. + + A curvilinear axis (a multi-dimensional grid array, produced by ``.map()`` + with ``space="conf"``) has no single 1-D coordinate array of its own to + search -- a coordinate value or slice string is resolved against a 1-D + cross-section instead (:func:`_curvilinear_coord_curve`), holding every + other axis of the same mapped block at cell 0, *unless* that axis was + already narrowed to a single cell by an earlier selector -- either in + this same call, or in a prior ``select()`` call in the chain (recorded + the only place it needs to be: the dataset's own values shape). Selecting + one axis of a block also narrows every sibling axis' grid array along + that same relative axis, so the block stays internally consistent and a + later selector on a sibling sees the narrowed cross-section rather than + the original full extent. If the array still varies along an + as-yet-unresolved sibling axis (a genuinely non-separable map, e.g. a + rotation), the coordinate/slice selector has no single answer and raises + -- pick an integer index for that sibling axis first (in the same call, + or an earlier one in the chain). + + Args: + data: Dataset whose point values are selected. + comp: Component selector such as ``"0"``, ``"0:3"``, or ``"0,2"``. + z0: Selector for coordinate direction 0. + z1: Selector for coordinate direction 1. + z2: Selector for coordinate direction 2. + z3: Selector for coordinate direction 3. + z4: Selector for coordinate direction 4. + z5: Selector for coordinate direction 5. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + + Raises: + ValueError: if ``data`` holds native modal DG coefficients (nodal/quad + value_forms of gkyl-backed data are point values and slice fine), + or a coordinate/slice selector targets a curvilinear axis whose + physical coordinate still varies along an unresolved sibling axis. + """ + if data.backend == "gkyl" and data.ctx.get("value_form", "modal") == "modal": + raise ValueError( + "select operates on interpolated (NumPy) values, or on gkyl-native " + "nodal/quad value_forms; call .interpolate()/.to_nodal()/" + ".to_quad() first -- slicing raw modal DG coefficients would mix " + "basis functions.") + zs = (z0, z1, z2, z3, z4, z5) + grid = list(data.grid) + values = data.values + num_dims = data.num_dims + values_idx = [slice(0, values.shape[d]) for d in range(num_dims + 1)] + + # ctx["mapped_axes"] records, for every dimension touched by a .map() + # call, the offset of the mapped block it belongs to; group them back + # into blocks so selecting one axis can keep every sibling's grid array + # (they all share the block's tensor shape) in sync. + mapped_axes = data.ctx.get("mapped_axes", {}) + block_dims: dict[int, list[int]] = {} + for dd, off in mapped_axes.items(): + block_dims.setdefault(off, []).append(dd) + # per-block set of relative axes already given a selector earlier in this + # same call -- lets a later axis's coordinate search skip the + # separability check on a sibling the caller has deliberately pinned. + touched: dict[int, set] = {} + + for d, z in enumerate(zs): + if d >= num_dims or z is None: + continue + grid_arr = grid[d] + curvilinear = grid_arr.ndim > 1 # a .map()-deformed grid axis + # a curvilinear array's own axis k corresponds to absolute dimension + # `offset + k` (map.py's mapped block), not to axis d of `grid` itself + # -- ctx["mapped_axes"] records each absolute dimension's block offset + # so the N-D array can be indexed on its own relative axis. + offset = mapped_axes.get(d, 0) + rel = d - offset if curvilinear else d + len_grid = grid_arr.shape[rel] if curvilinear else grid_arr.shape[0] + is_matching = values.shape[d] == len_grid + if curvilinear and isinstance(z, int): + idx = z + elif curvilinear: + coord_curve = _curvilinear_coord_curve(grid_arr, rel, d, offset, + values.shape, + touched.get(offset, set())) + idx = idx_parser(z, coord_curve, is_matching) + else: + idx = idx_parser(z, grid_arr, is_matching) + if isinstance(idx, int): + if idx < 0: + idx = values.shape[d] + idx + v_idx = slice(idx, idx + 1) + g_idx = slice(idx, idx + 1) if is_matching else slice(idx, idx + 2) + elif isinstance(idx, slice): + v_idx = idx + g_idx = idx if is_matching else slice(idx.start, idx.stop + 1) + else: + raise TypeError("Coordinate selector must be a single index or a slice.") + if curvilinear: + # every axis in the same mapped block shares this array shape, so + # slicing relative axis `rel` in lockstep keeps them all consistent + # -- a later selector on a sibling axis then resolves against the + # already-narrowed cross-section instead of the original full extent. + for dd in block_dims.get(offset, [d]): + arr = grid[dd] + grid[dd] = arr[tuple(g_idx if k == rel else slice(None) + for k in range(arr.ndim))] + touched.setdefault(offset, set()).add(rel) + else: + grid[d] = grid_arr[g_idx] + values_idx[d] = v_idx + + if comp is not None: + values_idx[-1] = idx_parser(comp) + + values_out = values[tuple(values_idx)] + if num_dims == values_out.ndim: # restore the squeezed component axis + values_out = values_out[..., np.newaxis] + + ctx_updates = {} + if data.backend == "gkyl": + # A nodal/quad value_form stays gkyl-native (REFACTOR_GKEYLL_FFI.md + # §3b): ``values`` above was only a read-only NumPy *view* of the native + # array for slicing purposes -- wrap the sliced result back into a + # native GkylArray so the dataset doesn't silently fall out of the gkyl + # backend (and lose its value_form) just for having been selected. + # Cell layout isn't derivable from the flat native array (see + # ``GDataState.set_values``), so it must be threaded through explicitly, + # the same way ``average``/``eval_at_coord_proj`` do. + ctx_updates["cells"] = np.array(values_out.shape[:-1], dtype=np.int64) + values_out = dg.rep.wrap(values_out) + + return data._result(grid, + values_out, + inplace=inplace, + tag=tag, + label=label, + **ctx_updates) diff --git a/src/postgkyl/operations/sort.py b/src/postgkyl/operations/sort.py new file mode 100644 index 00000000..7570cd24 --- /dev/null +++ b/src/postgkyl/operations/sort.py @@ -0,0 +1,33 @@ +"""The ``sort`` verb -- natural/numeric-order datasets by source filename.""" + +from __future__ import annotations + +from postgkyl import numerics +from postgkyl.gdatastate import flatten_datasets +from postgkyl.gdatastate.gdatastate import GDataState + + +def sort(*datasets: GDataState, reverse: bool = False) -> list[GDataState]: + """Reorder datasets by the natural/numeric sort of their source filename. + + Fixes the shell-glob/lexicographic-sort trap where ``field_10.gkyl`` sorts + before ``field_2.gkyl``: digit runs embedded in the filename are compared + as integers rather than character-by-character, so frame files come out in + increasing frame order regardless of digit-count padding (see + ``numerics.natural_sort_key``). + + Accepts ``sort(a, b)`` or ``sort([a, b])`` (flattened via + ``gdatastate.flatten_datasets``). No dataset is copied or mutated -- only + the returned list's order differs from the input. + + Args: + *datasets: the datasets to reorder, or lists/groups thereof. + reverse: sort in decreasing order instead of increasing. + + Returns: + The same datasets, reordered by their source filename's natural sort key. + """ + states = flatten_datasets(datasets) + return sorted(states, + key=lambda d: numerics.natural_sort_key(d.file_name), + reverse=reverse) diff --git a/src/postgkyl/operations/val2coord.py b/src/postgkyl/operations/val2coord.py new file mode 100644 index 00000000..205fc50e --- /dev/null +++ b/src/postgkyl/operations/val2coord.py @@ -0,0 +1,99 @@ +"""The ``val2coord`` verb -- build new datasets from columns of a DynVector.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from postgkyl.gdatastate.gdatastategroup import GDataStateGroup + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def _get_range(str_in: str, length: int) -> np.ndarray: + """Parse a comma list, a ``lo:hi[:step]`` slice, or a bare int into indices. + + Pure array/string logic, no ``GData`` coupling; kept local to this verb + (its grammar -- an optional step -- differs from ``numerics.idx_parser``'s + slice grammar, and it has no other caller). + """ + if len(str_in.split(",")) > 1: + return np.array(str_in.split(","), dtype=int) + elif str_in.find(":") >= 0: + parts = str_in.split(":") + s_idx = 0 if parts[0] == "" else int(parts[0]) + if s_idx < 0: + s_idx = length + s_idx + e_idx = length if parts[1] == "" else int(parts[1]) + if e_idx < 0: + e_idx = length + e_idx + inc = int(parts[2]) if len(parts) > 2 and parts[2] != "" else 1 + return np.arange(s_idx, e_idx, inc) + else: + return np.array([int(str_in)]) + + +def val2coord(data: "GDataState", + *, + x: str, + y: str, + periodic: bool = False, + tag: str | None = None, + label: str | None = None) -> GDataStateGroup: + """Build new (x, y) datasets from columns of a DynVector. + + Reinterprets columns of ``data`` (typically a DynVector / diagnostic + table) as plot-ready datasets: the ``x`` column(s) become the grid and the + ``y`` column(s) become the values. One output dataset is produced per + selected y-component. When more than one x-component is selected, their + count must match the number of y-components (paired one-to-one); a single + x-component is shared across all y-components. + + Args: + data: the source dataset whose last-axis columns are selected; must be + NumPy-backed. + x: component selector for the independent variable: an integer index, a + comma-separated list (e.g. ``'0,2'``), or a ``'lo:hi[:step]'`` slice. + y: component selector for the dependent variable(s); same forms as + ``x``. One output dataset is produced per selected y-component. + periodic: when True, append the first sample to the end of each output + (wrapping) so periodic data closes on itself. + tag: optional tag for the returned datasets. + label: optional label for the returned datasets. + + Returns: + A ``GDataStateGroup`` containing one dataset per selected y-component. + + Raises: + ValueError: if ``data`` is native modal (gkyl-backed), or if more than + one x-component is selected and their number does not equal the + number of y-components. + """ + if data.backend == "gkyl": + raise ValueError( + "val2coord operates on interpolated (NumPy) values; call .interpolate() " + "first -- raw DG coefficients are not tabular columns.") + values = data.values + x_comps = _get_range(x, values.shape[-1]) + y_comps = _get_range(y, values.shape[-1]) + + if len(x_comps) > 1 and len(x_comps) != len(y_comps): + raise ValueError( + f"val2coord: number of x-components ({len(x_comps):d}) is greater " + f"than 1 and not equal to the number of y-components " + f"({len(y_comps):d}).") + + out = [] + for i, yc in enumerate(y_comps): + xc = x_comps[i] if len(x_comps) > 1 else x_comps[0] + xv = values[..., xc] + yv = values[..., yc] + if periodic: + xv = np.append(xv, np.atleast_1d(xv[0]), axis=0) + yv = np.append(yv, np.atleast_1d(yv[0]), axis=0) + res = data._result([xv], yv[..., np.newaxis], tag=tag, label=label) + res.color = "C0" + out.append(res) + return GDataStateGroup(out) diff --git a/src/postgkyl/output/__init__.py b/src/postgkyl/output/__init__.py deleted file mode 100644 index 841c47a0..00000000 --- a/src/postgkyl/output/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Import plot -from .plot import plot - -from .plot import pgkyl_colorbar diff --git a/src/postgkyl/output/plot.py b/src/postgkyl/output/plot.py deleted file mode 100644 index f9b39427..00000000 --- a/src/postgkyl/output/plot.py +++ /dev/null @@ -1,677 +0,0 @@ -"""Module including custom Gkeyll plotting function""" -from __future__ import annotations - -from contextlib import nullcontext -from matplotlib import cm -from matplotlib import colors -from matplotlib import patches -from mpl_toolkits.axes_grid1 import make_axes_locatable -import mpl_toolkits.mplot3d # noqa: F401 (registers the '3d' projection) -from typing import Tuple, TYPE_CHECKING -import matplotlib as mpl -import matplotlib.axes -import matplotlib.figure -import matplotlib.pyplot as plt -import matplotlib.font_manager as fm -import numpy as np -import os.path - -from postgkyl.utils import input_parser -if TYPE_CHECKING: - from postgkyl import GData -# end - -# Helper functions -def pgkyl_colorbar(obj, fig : matplotlib.figure.Figure, cax : matplotlib.axes.Axes, - label: str = "", extend: bool | None = None): - divider = make_axes_locatable(cax) - cax2 = divider.append_axes("right", size="3%", pad=0.05) - return fig.colorbar(obj, cax=cax2, label=label or "", extend=extend) - -def get_xkcd_safely(xkcd=True): - # Target fonts Matplotlib looks for in xkcd mode - required_fonts = ['xkcd', 'xkcd Script', 'Comic Neue', 'Comic Sans MS'] - - # Check if at least one comic font is installed - available_fonts = [f.name for f in fm.fontManager.ttflist] - has_font = any(rf in available_fonts for rf in required_fonts) - - if not has_font: - print("\n" + "="*50) - print(" MINIMALIST XKCD FONT GUIDE") - print("="*50) - print("Matplotlib cannot find an XKCD-compatible font.") - print("\nTo fix this completely:") - print("1. Download the missing fonts (xkcd, xkcd-script, Comic Neue, or Comic Sans MS).") - print("2. Install the font(s) onto your operating system.") - print("3. Clear Matplotlib cache & restart your IDE:") - print(" -> import shutil, matplotlib as mpl") - print(" -> shutil.rmtree(mpl.get_cachedir())") - print("="*50 + "\n") - - # Fallback: Attempt standard sans-serif so the script keeps running - font_rc = {'font.family': 'sans-serif'} - else: - # Secure fallback to Comic Sans if original xkcd isn't there but Sans is - font_rc = {'font.family': 'Comic Sans MS'} - - return plt.xkcd, font_rc - - -def _get_nodal_grid(grid : list, cells: np.ndarray): - num_dims = len(grid) - grid_out = [] - if num_dims != len(cells): # sanity check - raise ValueError("Number dimensions for 'grid' and 'values' doesn't match") - # end - for d in range(num_dims): - if len(grid[d].shape) == 1: - if grid[d].shape[0] == cells[d]: - grid_out.append(grid[d]) - elif grid[d].shape[0] == cells[d] + 1: - grid_out.append(0.5 * (grid[d][:-1] + grid[d][1:])) - else: - raise ValueError("Something is terribly wrong...") - # end - else: - if grid[d].shape[d] == cells[d]: - grid_out.append(grid[d]) - elif grid[d].shape[d] == cells[d] + 1: - if num_dims == 1: - grid_out.append(0.5 * (grid[d][:-1] + grid[d][1:])) - else: - grid_out.append(0.5 * (grid[d][:-1, :-1] + grid[d][1:, 1:])) - # end - else: - raise ValueError("Something is terribly wrong...") - # end - # end - # end - return grid_out - - -def plot(data: GData | Tuple[list, np.ndarray], args: list = (), - figure: int | matplotlib.figure.Figure | str | None = None, - squeeze: bool = False, transpose: bool = False, num_axes: int = None, start_axes: int = 0, - num_subplot_row: int | None = None, num_subplot_col: int | None = None, - streamline: bool = False, sdensity: int = 1, - quiver: bool = False, - contour: bool = False, clevels: list | None = None, cnlevels: int | None = None, cont_label: bool = False, - surface: bool = False, comparison: bool = False, alpha: float | None = None, - diverging: bool = False, - lineouts: int | None = None, - xmin: float | None = None, xmax: float | None = None, xscale: float = 1.0, xshift: float = 0.0, - ymin: float | None = None, ymax: float | None = None, yscale: float = 1.0, yshift: float = 0.0, - zmin: float | None = None, zmax: float | None = None, zscale: float = 1.0, zshift: float = 0.0, - relax: bool = False, style: str | None = None, rcParams: dict | None = None, - legend: bool = True, label_prefix: str = "", colorbar: bool = True, - xlabel: str | None = None, ylabel: str | None = None, clabel: str | None = None, title: str | None = None, - subplot_titles: str | None = None, subplot_xlabels: str | None = None, subplot_ylabels: str | None = None, - logx: bool = False, logy: bool = False, logz: bool = False, - fixaspect: bool = False, aspect: float | None = None, - edgecolors: str | None = None, showgrid: bool = True, hashtag: bool = False, xkcd: bool = False, - color: str | None = None, markersize: float | None = None, - linewidth: float | None = None, linestyle: float | None = None, - figsize: tuple | None = None, - jet: bool = False, cmap: str | None = None, - cval: float | None = None, cval_min: float | None = None, cval_max: float | None = None, - **kwargs): - """Plots Gkeyll data. - - Unifies the plotting across a wide range of Gkyl applications. Can - be used for both 1D an 2D data. Uses a proper colormap by default. - - For 1D data, passing ``cmap`` together with ``cval`` colors the line by mapping - ``cval`` onto the colormap. ``cval_min``/``cval_max`` set the normalization - range (typically the min/max of the ``cval`` values across all curves), so that - several curves drawn into the same axes share a consistent color scale. - - For 2D data, ``surface`` draws a 3D surface instead of a ``pcolormesh``. When - several 2D datasets are overlaid onto the same axes for comparison, set - ``comparison`` so that each surface/contour gets a distinct color and a legend - entry (instead of overlapping and hiding each other). ``alpha`` controls the - surface transparency. - """ - - # Default to Postgkyl style file file if no style is specified - # Use the rcParams dictionary which is passed with click contex - if bool(style): - plt.style.use(style) - elif bool(rcParams): - for key in rcParams: - mpl.rcParams[key] = rcParams[key] - # end - else: - plt.style.use(f"{os.path.dirname(os.path.realpath(__file__)):s}/postgkyl.mplstyle") - # end - - # Process input parameters - if not bool(aspect): - aspect = 1.0 - # end - - if bool(cmap): - mpl.rcParams["image.cmap"] = cmap - elif bool(diverging): - mpl.rcParams["image.cmap"] = "RdBu_r" - # end - - # This should not be used on its own; however, it can be useful for - # comparing results with literature - if bool(jet): - mpl.rcParams["image.cmap"] = "jet" - # end - - # The most important thing - if xkcd: - xkcd_cm, xkcd_rc = get_xkcd_safely() - else: - xkcd_cm = nullcontext - xkcd_rc = {} - # end - - with xkcd_cm(), mpl.rc_context(rc=xkcd_rc): - - if not bool(color) and not isinstance(data, tuple): - cl = data.color - # end - if bool(color): - mpl.rcParams["lines.color"] = color - # end - if bool(linewidth): - mpl.rcParams["lines.linewidth"] = linewidth - # end - if bool(linestyle): - mpl.rcParams["lines.linestyle"] = linestyle - # end - - # Get the handles on the grid and values - grid_in, values = input_parser(data) - grid = grid_in.copy() - - if isinstance(data, tuple): - if len(grid) == len(values.shape): - num_dims = len(values.squeeze().shape) - else: - num_dims = len(values[..., 0].squeeze().shape) - # end - lg = len(grid) - lower, upper, cells = np.zeros(lg), np.zeros(lg), np.zeros(lg) - for d in range(lg): - lower[d] = np.min(grid[d]) - upper[d] = np.max(grid[d]) - if len(grid[d].shape) == 1: - cells[d] = len(grid[d]) - else: - cells[d] = len(grid[d][d]) - # end - # end - else: # GData - num_dims = data.get_num_dims(squeeze=True) - lower, upper = data.get_bounds() - cells = data.get_num_cells() - # end - if num_dims > 2: - raise ValueError("Only 1D and 2D plots are currently supported") - # end - - # Squeeze the data (get rid of "collapsed" dimensions) - axes_labels = ["$z_0$", "$z_1$", "$z_2$", "$z_3$", "$z_4$", "$z_5$"] - if len(grid) > num_dims: - idx = [] - for dim, g in enumerate(grid): - if cells[dim] <= 1: - idx.append(dim) - # end - grid[dim] = g.squeeze() - # end - if bool(idx): - for i in reversed(idx): - grid.pop(i) - # end - lower = np.delete(lower, idx) - upper = np.delete(upper, idx) - cells = np.delete(cells, idx) - axes_labels = np.delete(axes_labels, idx) - values = np.squeeze(values, tuple(idx)) - - # c2p grids - if len(grid[0].shape) > 1: - for d in range(num_dims): - for i in reversed(idx): - grid[d] = np.mean(grid[d], axis=i) - # end - # end - # end - # end - # end - - # Swap the horizontal and vertical axes. - if transpose and num_dims == 2: - values = np.swapaxes(values, 0, 1) - g0, g1 = grid[1], grid[0] - if g0.ndim > 1: - g0, g1 = g0.transpose(), g1.transpose() - # end - grid[0], grid[1] = g0, g1 - lower[0], lower[1] = lower[1], lower[0] - upper[0], upper[1] = upper[1], upper[0] - cells[0], cells[1] = cells[1], cells[0] - axes_labels[0], axes_labels[1] = axes_labels[1], axes_labels[0] - # end - - # Get the number of components and an indexer - step = 2 if bool(streamline or quiver) else 1 - num_comps = values.shape[-1] - idx_comps = range(int(np.floor(num_comps / step))) - if num_axes: - num_comps = num_axes - else: - num_comps = len(idx_comps) - # end - - # Create axis labels - if xlabel is None: - xlabel = axes_labels[0] if lineouts != 1 else axes_labels[1] - if xshift != 0.0 and xscale != 1.0: - xlabel = rf"({xlabel:s} + {xshift:.2e}) $\times$ {xscale:.2e}" - elif xshift != 0.0: - xlabel = rf"{xlabel:s} + {xshift:.2e}" - elif xscale != 1.0: - xlabel = rf"{xlabel:s} $\times$ {xscale:.2e}" - # end - # end - if ylabel is None and num_dims == 2 and lineouts is None: - ylabel = axes_labels[1] - if yshift != 0.0 and yscale != 1.0: - ylabel = rf"({ylabel:s} + {yshift:.2e}) $\times$ {yscale:.2e}" - elif xshift != 0.0: - ylabel = rf"{ylabel:s} + {yshift:.2e}" - elif xscale != 1.0: - ylabel = rf"{ylabel:s} $\times$ {yscale:.2e}" - # end - # end - if zscale != 1.0: - if clabel: - clabel = rf"{clabel:s} $\times$ {zscale:.3e}" - else: - clabel = rf"$\times$ {zscale:.3e}" - # end - # end - - if transpose and num_dims == 1: - xlabel, ylabel = ylabel, xlabel - # end - - # Surface plots need 3D axes; only meaningful for 2D data. - use_3d = bool(surface) and num_dims == 2 - subplot_kw = {"projection": "3d"} if use_3d else {} - - if bool(figsize): - figsize = (int(figsize.split(",")[0]), int(figsize.split(",")[1])) - # end - if figure is None: - fig = plt.figure(figsize=figsize) - elif isinstance(figure, int): - fig = plt.figure(figure, figsize=figsize) - elif isinstance(figure, matplotlib.figure.Figure): - fig = figure - elif isinstance(figure, str): - fig = plt.figure(int(figure), figsize=figsize) - else: - raise TypeError( - "'fig' keyword needs to be one of " "None (default), int, or MPL Figure" - ) - # end - - # Axes - if fig.axes: - ax = fig.axes - if squeeze is False and num_comps > len(ax): - raise ValueError("Trying to plot into figure with not enough axes") - # end - else: - if squeeze: # Plotting into 1 panel - fig.subplots(1, 1, subplot_kw=subplot_kw) - ax = fig.axes - ax[0].set_xlabel(xlabel) - ax[0].set_ylabel(ylabel) - if title is not None: - ax[0].set_title(title, y=1.08) - # end - else: # Plotting each components into its own subplot - if num_subplot_row is not None: - num_rows = num_subplot_row - num_cols = int(np.ceil(num_comps/num_rows)) - elif num_subplot_col is not None: - num_cols = num_subplot_col - num_rows = int(np.ceil(num_comps/num_cols)) - else: - sr = np.sqrt(num_comps) - if sr == np.ceil(sr): - num_rows = int(sr) - num_cols = int(sr) - elif np.ceil(sr) * np.floor(sr) >= num_comps: - num_rows = int(np.floor(sr)) - num_cols = int(np.ceil(sr)) - else: - num_rows = int(np.ceil(sr)) - num_cols = int(np.ceil(sr)) - # end - # end - - if num_dims == 1 or lineouts is not None: - fig.subplots(num_rows, num_cols, sharex=True, subplot_kw=subplot_kw) - elif use_3d: # 3D axes cannot share x/y with each other - fig.subplots(num_rows, num_cols, subplot_kw=subplot_kw) - else: # In 2D, share y-axis as well - fig.subplots(num_rows, num_cols, sharex=True, sharey=True) - # end - ax = fig.axes - # Removing extra axes - for i in range(num_comps, len(ax)): - ax[i].axis("off") - # end - # Add labels as super labels and titles - if bool(title): - fig.suptitle(title) - if bool(xlabel): - fig.supxlabel(xlabel) - if bool(ylabel): - fig.supylabel(ylabel) - - for ax_idx, _ in enumerate(ax): - if bool(subplot_titles): - title = subplot_titles.split(",")[ax_idx] if ax_idx < len(subplot_titles.split(",")) else "" - else: - title = "" - # end - if bool(subplot_xlabels): - xlabel = subplot_xlabels.split(",")[ax_idx] if ax_idx < len(subplot_xlabels.split(",")) else "" - else: - xlabel = "" - # end - if bool(subplot_ylabels): - ylabel = subplot_ylabels.split(",")[ax_idx] if ax_idx < len(subplot_ylabels.split(",")) else "" - else: - ylabel = "" - # end - - ax[ax_idx].set_xlabel(xlabel) - ax[ax_idx].set_ylabel(ylabel) - if bool(title): - ax[ax_idx].set_title(title, y=1.08) - # end - # end - # end - # end - - for comp in idx_comps: - cax = ax[0] if squeeze else ax[comp + start_axes] - label = f"{label_prefix:s}_c{comp:d}".strip("_") if len(idx_comps) > 1 else label_prefix - - if num_dims == 1: - nodal_grid = _get_nodal_grid(grid, cells) - x = (nodal_grid[0] + xshift)*xscale - y = (values[..., comp] + yshift)*yscale - if transpose: # put the coordinate on the vertical axis - x, y = y, x - # end - # Color the line from the colormap when a 'cval' is given (1D only). - line_color = color - if bool(cmap) and cval is not None: - if cval_max is not None and cval_min is not None and cval_max != cval_min: - t = (cval - cval_min)/(cval_max - cval_min) - else: - t = 0.5 - # end - line_color = plt.get_cmap(cmap)(t) - # end - im = cax.plot(x, y, *args, color=line_color, label=label, markersize=markersize) - # Add a colorbar describing the cval-to-color mapping once per axes. - if (bool(cmap) and cval is not None and colorbar - and cval_max is not None and cval_min is not None and cval_max != cval_min - and not getattr(cax, "_pgkyl_cval_cbar", False)): - mappable = cm.ScalarMappable( - norm=colors.Normalize(vmin=cval_min, vmax=cval_max), cmap=plt.get_cmap(cmap)) - pgkyl_colorbar(mappable, fig, cax, label=clabel) - cax._pgkyl_cval_cbar = True - # end - - elif num_dims == 2: - extend = None - - if surface: - nodal_grid = _get_nodal_grid(grid, cells) - xg = (nodal_grid[0] + xshift)*xscale - yg = (nodal_grid[1] + yshift)*yscale - z = (values[..., comp].transpose() + zshift)*zscale - if xg.ndim == 1: - xg, yg = np.meshgrid(xg, yg) - else: - xg, yg = xg.transpose(), yg.transpose() - # end - # Count how many overlays already live on these axes so each gets a - # distinct color (used for both surface and contour comparisons). - count = getattr(cax, "_pgkyl_overlay_count", 0) - cax._pgkyl_overlay_count = count + 1 - if comparison or bool(color): - surf_color = color if bool(color) else f"C{count:d}" - im = cax.plot_surface(xg, yg, z, color=surf_color, - alpha=alpha if alpha is not None else 0.6, - linewidth=0, antialiased=True, shade=True) - if label: - handles = getattr(cax, "_pgkyl_handles", []) - handles.append(patches.Patch(color=surf_color, label=label)) - cax._pgkyl_handles = handles - # end - else: - im = cax.plot_surface(xg, yg, z, cmap=mpl.rcParams["image.cmap"], - alpha=alpha if alpha is not None else 1.0, - linewidth=0, antialiased=True) - if colorbar: - fig.colorbar(im, ax=cax, label=clabel or "", shrink=0.6, pad=0.1) - # end - # end - if clabel: - cax.set_zlabel(clabel) - # end - if zmin is not None or zmax is not None: - cax.set_zlim(zmin, zmax) - # end - colorbar = False - - elif contour: - levels = 10 - if cnlevels: - levels = int(cnlevels) - 1 - elif clevels: - if ":" in clevels: - s = clevels.split(":") - levels = np.linspace(float(s[0]), float(s[1]), int(s[2])) - else: - levels = np.array(clevels.split(",")) - # Filter out empty elements - levels = np.array(list(filter(None, levels))) - # end - # end - if isinstance(levels, np.ndarray) and len(levels) == 1: - colorbar = False - # end - nodal_grid = _get_nodal_grid(grid, cells) - x = (nodal_grid[0] + xshift) * xscale - y = (nodal_grid[1] + yshift) * yscale - z = (values[..., comp].transpose() + zshift) * zscale - cont_colors = color - if comparison and not bool(color): - # Give each overlaid dataset a distinct, single color + legend entry. - count = getattr(cax, "_pgkyl_overlay_count", 0) - cax._pgkyl_overlay_count = count + 1 - cont_colors = f"C{count:d}" - if label: - handles = getattr(cax, "_pgkyl_handles", []) - handles.append(patches.Patch(color=cont_colors, label=label)) - cax._pgkyl_handles = handles - # end - colorbar = False - # end - im = cax.contour(x, y, z, levels, *args, origin="lower", colors=cont_colors, linewidths=linewidth) - if cont_label: - cax.clabel(im, inline=1) - # end - - elif quiver: - skip = int(np.max((len(grid[0]), len(grid[1])))//15) - skip2 = int(skip//2) - nodal_grid = _get_nodal_grid(grid, cells) - if len(nodal_grid[0].shape) == 1: - x = (nodal_grid[0][skip2::skip] + xshift)*xscale - y = (nodal_grid[1][skip2::skip] + yshift)*yscale - else: - x = (nodal_grid[0][skip2::skip, skip2::skip] + xshift)*xscale - y = (nodal_grid[1][skip2::skip, skip2::skip] + yshift)*yscale - # end - z1 = (values[skip2::skip, skip2::skip, 2 * comp].transpose() + zshift)*zscale - z2 = (values[skip2::skip, skip2::skip, 2 * comp + 1].transpose() + zshift)*zscale - im = cax.quiver(x, y, z1, z2) - - elif streamline: - if bool(color): - cl = color - else: - # magnitude - cl = np.sqrt( - values[..., 2 * comp]**2 + values[..., 2 * comp + 1]**2 - ).transpose() - # end - nodal_grid = _get_nodal_grid(grid, cells) - x = (nodal_grid[0] + xshift)*xscale - y = (nodal_grid[1] + yshift)*yscale - z1 = (values[..., 2 * comp].transpose() + zshift)*zscale - z2 = (values[..., 2 * comp + 1].transpose() + zshift)*zscale - im = cax.streamplot(x, y, z1, z2, *args, - density=sdensity, broken_streamlines=False, color=cl, linewidth=linewidth) - - elif lineouts is not None: - num_lines = values.shape[1] if lineouts == 0 else values.shape[0] - nodal_grid = _get_nodal_grid(grid, cells) - - if lineouts == 0: - x = (nodal_grid[0] + xshift)*xscale - vmin = (nodal_grid[1][0] + yshift)*yscale - vmax = (nodal_grid[1][-1] + yshift)*yscale - label = clabel or axes_labels[1] - else: - x = (nodal_grid[1] + xshift)*xscale - vmin = (nodal_grid[0][0] + yshift)*yscale - vmax = (nodal_grid[0][-1] + yshift)*yscale - label = clabel or axes_labels[0] - # end - idx = [slice(0, u) for u in values.shape] - idx[-1] = comp - for line in range(num_lines): - color = cm.inferno(line / (num_lines - 1)) - if lineouts == 0: - idx[1] = line - else: - idx[0] = line - # end - y = (values[tuple(idx)] + yshift)*yscale - im = cax.plot(x, y, *args, color=color) - # end - mappable = cm.ScalarMappable( - norm=colors.Normalize(vmin=vmin, vmax=vmax, clip=False), cmap=cm.inferno - ) - pgkyl_colorbar(mappable, fig, cax, label=label) - colorbar = False - legend = False - - else: - if zmin is not None and zmax is not None: - extend = "both" - elif zmax is not None: - extend = "max" - elif zmin is not None: - extend = "min" - # end - x = (grid[0] + xshift)*xscale - y = (grid[1] + yshift)*yscale - z = (values[..., comp].transpose() + zshift)*zscale - if len(x) == z.shape[1] or len(y) == z.shape[0]: - nodal_grid = _get_nodal_grid(grid, cells) - x = (nodal_grid[0] + xshift)*xscale - y = (nodal_grid[1] + yshift)*yscale - # end - if len(x.shape) > 1: - x, y = x.transpose(), y.transpose() - # end - if diverging: - zmax = np.abs(z).max() - zmin = -zmax - # end - vmax, vmin = zmax, zmin - norm = None - if logz: - if diverging: - tmp = vmax/1000 - norm = colors.SymLogNorm( - linthresh=tmp, linscale=tmp, vmin=vmin, vmax=vmax, base=10 - ) - else: - norm = colors.LogNorm(vmin=vmin, vmax=vmax) - # end - vmin, vmax = None, None - # end - im = cax.pcolormesh(x, y, z, - norm=norm, vmin=vmin, vmax=vmax, edgecolors=edgecolors, - linewidth=0.1, shading="auto", *args) - # end - if not bool(color) and colorbar and not streamline: - pgkyl_colorbar(im, fig, cax, extend=extend, label=clabel) - # end - else: - raise ValueError(f"{num_dims:d}D data not supported") - # end - - cax.grid(showgrid) - # Legend - if legend: - if getattr(cax, "_pgkyl_handles", None): - # Overlaid 2D datasets (surface/contour comparison): real legend. - cax.legend(handles=cax._pgkyl_handles, loc=0) - elif num_dims == 1 and label != "": - cax.legend(loc=0) - elif not (surface and num_dims == 2): - cax.text(0.03, 0.96, label, - bbox={"facecolor": "w", "edgecolor": "w", "alpha": 0.8, "boxstyle": "round"}, - verticalalignment="top", horizontalalignment="left", transform=cax.transAxes) - # end - # end - if hashtag: - cax.text(0.97, 0.03, "#pgkyl", - bbox={"facecolor": "w", "edgecolor": "w", "alpha": 0.8, "boxstyle": "round"}, - verticalalignment="bottom", horizontalalignment="right", transform=cax.transAxes) - # end - if logx: - cax.set_xscale("log") - # end - if logy: - cax.set_yscale("log") - # end - if num_dims == 1 and not relax: # this causes troubles with contours - plt.autoscale(enable=True, axis="x", tight=True) - plt.autoscale(enable=True, axis="y") - # end - if xmin is not None or xmax is not None: - cax.set_xlim(xmin, xmax) - # end - if ymin is not None or ymax is not None: - cax.set_ylim(ymin, ymax) - # end - if fixaspect and not (surface and num_dims == 2): - plt.setp(cax, aspect=aspect) - # end - # end - - plt.tight_layout() - return im diff --git a/src/postgkyl/pgkyl.py b/src/postgkyl/pgkyl.py deleted file mode 100755 index 7b065457..00000000 --- a/src/postgkyl/pgkyl.py +++ /dev/null @@ -1,185 +0,0 @@ -#!/usr/bin/env python3 -"""Command line entry point for postgkyl. - -Uses click (https://click.palletsprojects.com/en) commands to wrap pgkyl functions. -""" - -from glob import glob -import click -import os.path -import sys -import time - -from postgkyl import __version__ -from postgkyl.commands import DataSpace -from postgkyl.utils import load_style, verb_print -import postgkyl.commands as cmd - - -def _print_version(ctx, param, value): - if not value or ctx.resilient_parsing: - return - # end - click.echo(f"Postgkyl {__version__} ({sys.platform})") - click.echo(f"Python version: {sys.version}".format()) - click.echo("Copyright 2016-2024 Gkeyll Team") - click.echo("Postgkyl can be used freely for research at universities,") - click.echo("national laboratories, and other non-profit institutions.") - click.echo("There is NO warranty.\n") - click.echo("Spam, egg, sausage, and spam.") - ctx.exit() - - -class PgkylCommandGroup(click.Group): - """Custom pgkyl click command group class. - - It allows to: - - use shortened versions of command names - - use a file name as a command - """ - - def get_command(self, ctx, cmd_name): - # cmd_name is a full name of a pgkyl command - rv = click.Group.get_command(self, ctx, cmd_name) - if rv is not None: - return rv - # end - - # cmd_name is an abreviation of a pgkyl command - matches = [x for x in self.list_commands(ctx) if x.startswith(cmd_name)] - if matches and len(matches) == 1: - return click.Group.get_command(self, ctx, matches[0]) - elif matches: - ctx.fail(f"Too many matches for '{cmd_name}': {', '.join(sorted(matches))}") - # end - - # cmd_name is a data set - if glob(cmd_name): - ctx.obj["in_data_strings"].append(cmd_name) - return click.Group.get_command(self, ctx, "load") - # end - - ctx.fail(f"'{cmd_name}' does not match either command name nor a data file") - - -# The command line mode entry command -@click.command(name="pgkyl", cls=PgkylCommandGroup, chain=True, - context_settings=dict(help_option_names=["-h", "--help"])) -@click.option("--verbose", "-v", is_flag=True, help="Turn on verbosity.") -@click.option("--batch-mode", is_flag=True, help="Run in batch mode (no plots will be shown).") -@click.option("--saveframes-prefix", default=os.path.expanduser("~")+"/pg", - help="Output prefix to use for plot output in batch mode.") -@click.option("--version", is_flag=True, callback=_print_version, expose_value=False, - is_eager=True, help="Print the version information.") -@click.option("--z0", help="Partial file load: 0th coord (either int or slice)") -@click.option("--z1", help="Partial file load: 1st coord (either int or slice)") -@click.option("--z2", help="Partial file load: 2nd coord (either int or slice)") -@click.option("--z3", help="Partial file load: 3rd coord (either int or slice)") -@click.option("--z4", help="Partial file load: 4th coord (either int or slice)") -@click.option("--z5", help="Partial file load: 5th coord (either int or slice)") -@click.option("--component", "-c", help="Partial file load: comps (either int or slice)") -@click.option("--compgrid", is_flag=True, help="Disregard the mapped grid information") -@click.option("--varname", "-d", multiple=True, - help="Specify the Adios variable name (default is 'CartGridField')") -@click.option("--c2p", help="Specify the file name containing c2p mapped coordinates") -@click.option("--c2p-vel", "c2p_vel", - help="Specify the file name containing c2p mapped velocity coordinates") -@click.option("--style", help="Sets Maplotlib rcParams style file.") -@click.pass_context -def cli(ctx, **kwargs): - """Postprocessing and plotting tool for Gkeyll data. - - Datasets can be loaded, processed and plotted using a command chaining mechanism. For - full documentation see the Gkeyll documentation webpages - (https://gkeyll.readthedocs.io). Help for individual commands can be obtained using - the --help option for that command. - """ - ctx.obj = {} # The main contex object - ctx.obj["start_time"] = time.time() # Timings are written in the verbose mode - if kwargs["verbose"]: - ctx.obj["verbose"] = True - verb_print(ctx, "This is Postgkyl running in verbose mode.") - else: - ctx.obj["verbose"] = False - # end - - ctx.obj["batch_mode"] = False - if kwargs["batch_mode"]: - ctx.obj["batch_mode"] = True - #end - - ctx.obj["saveframes_prefix"] = kwargs["saveframes_prefix"] - - ctx.obj["in_data_strings"] = [] - ctx.obj["in_data_strings_loaded"] = 0 - - ctx.obj["data"] = DataSpace() - - ctx.obj["fig"] = "" - ctx.obj["ax"] = "" - - ctx.obj["compgrid"] = kwargs["compgrid"] - ctx.obj["global_var_names"] = kwargs["varname"] - ctx.obj["global_cuts"] = (kwargs["z0"], kwargs["z1"], kwargs["z2"], - kwargs["z3"], kwargs["z4"], kwargs["z5"], kwargs["component"]) - ctx.obj["global_c2p"] = kwargs["c2p"] - ctx.obj["global_c2p_vel"] = kwargs["c2p_vel"] - - ctx.obj["rcParams"] = {} - fn = kwargs["style"] if kwargs["style"] else f"{os.path.dirname(os.path.realpath(__file__))}/output/postgkyl.mplstyle" - load_style(ctx, fn) - - -# Hook the individual commands into pgkyl -cli.add_command(cmd.config) -cli.add_command(cmd.activate) -cli.add_command(cmd.agyro) -cli.add_command(cmd.mom_agyro) -cli.add_command(cmd.animate) -cli.add_command(cmd.collect) -cli.add_command(cmd.current) -cli.add_command(cmd.deactivate) -cli.add_command(cmd.dg_evproj) -cli.add_command(cmd.dg_avg) -cli.add_command(cmd.differentiate) -cli.add_command(cmd.energetics) -cli.add_command(cmd.euler) -cli.add_command(cmd.mhd) -cli.add_command(cmd.ev) -cli.add_command(cmd.extractinput) -cli.add_command(cmd.fft) -cli.add_command(cmd.gk_nodes) -cli.add_command(cmd.dg_local_poly) -cli.add_command(cmd.gk_distf) -cli.add_command(cmd.gk_load_quantity) -cli.add_command(cmd.grid) -cli.add_command(cmd.growth) -cli.add_command(cmd.info) -cli.add_command(cmd.integrate) -cli.add_command(cmd.interpolate) -cli.add_command(cmd.laguerrecompose) -cli.add_command(cmd.listoutputs) -cli.add_command(cmd.load) -cli.add_command(cmd.magsq) -cli.add_command(cmd.mask) -cli.add_command(cmd.gk_energy_balance) -cli.add_command(cmd.gk_particle_balance) -cli.add_command(cmd.gk_rz) -cli.add_command(cmd.gk_fluxsurf) -cli.add_command(cmd.plot) -cli.add_command(cmd.pr) -cli.add_command(cmd.relchange) -cli.add_command(cmd.select) -cli.add_command(cmd.style) -cli.add_command(cmd.tenmoment) -cli.add_command(cmd.trajectory) -cli.add_command(cmd.val2coord) -cli.add_command(cmd.velocity) -cli.add_command(cmd.write) -cli.add_command(cmd.transformframe) -cli.add_command(cmd.pkpm) - -if __name__ == "__main__": - ctx = [] - cli(ctx) -# end diff --git a/src/postgkyl/py.typed b/src/postgkyl/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/src/postgkyl/render/__init__.py b/src/postgkyl/render/__init__.py new file mode 100644 index 00000000..24331156 --- /dev/null +++ b/src/postgkyl/render/__init__.py @@ -0,0 +1,19 @@ +"""Visualization backends (a backend layer used by the fluent surface). + +Re-exporting backend functions here shadows the corresponding submodule +references the import machinery would otherwise set on this package. +``render.animate`` and ``render.plotly`` therefore resolve to the canonical +functions, matching ``render.plot`` and ``render.pyvista``. Other public +Plotly names are re-exported alongside ``plotly`` for the same reason. +""" + +from . import labels, style +from .animate import animate +from .matplotlib import plot +from .plotly import open_preview, plotly, plotly_animate, save_rotating_plotly_figure +from .pyvista import pyvista + +__all__ = [ + "plot", "animate", "labels", "style", "plotly", "plotly_animate", + "save_rotating_plotly_figure", "open_preview", "pyvista" +] diff --git a/src/postgkyl/render/_ffmpeg.py b/src/postgkyl/render/_ffmpeg.py new file mode 100644 index 00000000..353095d5 --- /dev/null +++ b/src/postgkyl/render/_ffmpeg.py @@ -0,0 +1,31 @@ +"""Shared ffmpeg discovery, used by both ``animate.py`` and ``plotly.py``. + +PATH is checked first, so a system or conda-provided ``ffmpeg`` (the +``environment.yml`` route) always wins; falling back to the bundled +``imageio-ffmpeg`` static binary means a bare ``pip install`` also gets a +working ffmpeg, with no system package manager and no conda involved. +""" + +from __future__ import annotations + +import shutil + + +def resolve_ffmpeg() -> str | None: + path = shutil.which("ffmpeg") + if path is not None: + return path + try: + import imageio_ffmpeg + except ImportError: + return None + return imageio_ffmpeg.get_ffmpeg_exe() + + +def require_ffmpeg(context: str) -> str: + path = resolve_ffmpeg() + if path is None: + raise RuntimeError( + f"{context}: ffmpeg is required but was not found on PATH and " + "imageio-ffmpeg is not installed.") + return path diff --git a/src/postgkyl/render/_prep.py b/src/postgkyl/render/_prep.py new file mode 100644 index 00000000..c7435cbd --- /dev/null +++ b/src/postgkyl/render/_prep.py @@ -0,0 +1,186 @@ +"""Dataset -> plottable-array preparation, shared by every render backend. + +Private to ``render/``: this is the one concern the old tree split across +``utils/load_plot_data.py`` (dataset -> grid/values/dimensionality) and +``utils/axis_and_grid_prep.py`` (squeeze collapsed axes, resolve axis/colorbar +label defaults). Here it collapses to a single function over +:class:`~postgkyl.gdatastate.gdatastate.GDataState` -- the new container already exposes +``grid``/``values``/``num_dims`` uniformly, so there is no dual "GData or +tuple" input to dispatch on (contrast the old ``load_plot_data``). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from postgkyl.gdatastate.gdatastate import GDataState + + +def default_axis_labels(num_dims: int) -> list[str]: + """Default per-axis labels ``$z_0$``, ``$z_1$``, ... (mathtext).""" + return [rf"$z_{i}$" for i in range(num_dims)] + + +def format_axis_label(label: str, shift: float, scale: float) -> str: + """Annotate an axis label with its shift/scale, matching the old style.""" + if shift != 0.0 and scale != 1.0: + return rf"({label:s} + {shift:.2e}) $\times$ {scale:.2e}" + if shift != 0.0: + return rf"{label:s} + {shift:.2e}" + if scale != 1.0: + return rf"{label:s} $\times$ {scale:.2e}" + return label + + +def squeeze_collapsed_axes( + grid: list[np.ndarray], + values: np.ndarray) -> tuple[list[np.ndarray], np.ndarray]: + """Drop grid axes with exactly one cell (e.g. a ``select()``-ed coordinate). + + Curvilinear (multi-dimensional, ``.map()``-produced) coordinate arrays + cannot simply be indexed on the dropped axis -- every coordinate array + spans all dimensions jointly -- so each is averaged along it first (a + size-1 axis is unaffected by the mean); the now-redundant axis entry is + then removed from the coordinate list. + + Args: + grid: One nodal (edge) coordinate array per dimension. + values: Cell values, shape ``(*cells, num_comps)``. + + Returns: + ``(grid, values)`` with every size-1 axis removed. + """ + num_dims = len(grid) + cells = values.shape[:num_dims] + drop = [d for d in range(num_dims) if cells[d] <= 1] + if not drop: + return list(grid), values + + grid = [np.asarray(g) for g in grid] + if any(g.ndim > 1 for g in grid): + for d in range(num_dims): + for i in reversed(drop): + grid[d] = np.mean(grid[d], axis=i) + for i in reversed(drop): + grid.pop(i) + values = np.squeeze(values, tuple(drop)) + return grid, values + + +def subplot_grid(num_comps: int, + num_rows: int | None = None, + num_cols: int | None = None) -> tuple[int, int]: + """Choose a near-square ``(rows, cols)`` layout for ``num_comps`` panels.""" + if num_rows is not None: + return num_rows, int(np.ceil(num_comps / num_rows)) + if num_cols is not None: + return int(np.ceil(num_comps / num_cols)), num_cols + sr = np.sqrt(num_comps) + if sr == np.ceil(sr): + return int(sr), int(sr) + if np.ceil(sr) * np.floor(sr) >= num_comps: + return int(np.floor(sr)), int(np.ceil(sr)) + return int(np.ceil(sr)), int(np.ceil(sr)) + + +@dataclass(frozen=True) +class PlotPanel: + """Squeezed, label-resolved view of one dataset, ready for a render call.""" + grid: list[np.ndarray] + values: np.ndarray + num_dims: int + num_comps: int + xlabel: str + ylabel: str + clabel: str + + +def resolve_axis_labels(*, + xlabel: str | None, + ylabel: str | None, + zlabel: str | None, + clabel: str, + num_dims: int, + xshift: float = 0.0, + yshift: float = 0.0, + zshift: float = 0.0, + xscale: float = 1.0, + yscale: float = 1.0, + zscale: float = 1.0) -> tuple[str, str, str, str]: + """Infer default ``$z_i$`` labels and apply shift/scale annotations. + + Shared by the 2-D (``matplotlib``, no real ``z`` axis) and 3-D + (``plotly``, ``z`` is a genuine coordinate) backends: with ``num_dims`` + dimensions, defaults are ``z_0..z_{num_dims-1}`` distributed across + ``xlabel``/``ylabel``/``zlabel`` in that order (only as many as apply). + """ + labels = default_axis_labels(max(num_dims, 3)) + if xlabel is None: + xlabel = labels[0] if num_dims > 0 else "" + if ylabel is None: + ylabel = labels[1] if num_dims > 1 else "" + if zlabel is None: + zlabel = labels[2] if num_dims > 2 else labels[-1] + xlabel = format_axis_label(xlabel, xshift, xscale) + ylabel = format_axis_label(ylabel, yshift, yscale) + zlabel = format_axis_label(zlabel, zshift, zscale) + if zscale != 1.0: + clabel = (rf"{clabel:s} $\times$ {zscale:.3e}" + if clabel else rf"$\times$ {zscale:.3e}") + return xlabel, ylabel, zlabel, clabel + + +def prep_plot_data(data: "GDataState", + *, + xlabel: str | None = None, + ylabel: str | None = None, + clabel: str = "", + xshift: float = 0.0, + yshift: float = 0.0, + zshift: float = 0.0, + xscale: float = 1.0, + yscale: float = 1.0, + zscale: float = 1.0) -> PlotPanel: + """Squeeze collapsed axes and resolve axis/colorbar label defaults. + + Args: + data: The dataset to prepare (point-value/NumPy-backed; the ``plot`` + verb has already bridged any modal data through its NumPy shadow). + xlabel: Explicit x-axis label; auto-derived (``$z_0$``) when ``None``. + ylabel: Explicit y-axis label; auto-derived (``$z_1$``) when ``None`` + and the (squeezed) dataset is 2-D, else empty. + clabel: Colorbar label base text; annotated with ``zscale`` when it is + not 1. + xshift, yshift, zshift: Additive shifts recorded in the axis labels + (the caller applies them to the plotted arrays). + xscale, yscale, zscale: Multiplicative scales recorded in the axis + labels (the caller applies them to the plotted arrays). + + Returns: + A :class:`PlotPanel` with the squeezed grid/values and resolved labels. + """ + grid, values = squeeze_collapsed_axes(list(data.grid), data.values) + num_dims = len(grid) + xlabel, ylabel, _zlabel, clabel = resolve_axis_labels(xlabel=xlabel, + ylabel=ylabel, + zlabel="", + clabel=clabel, + num_dims=num_dims, + xshift=xshift, + yshift=yshift, + zshift=zshift, + xscale=xscale, + yscale=yscale, + zscale=zscale) + + return PlotPanel(grid=grid, + values=values, + num_dims=num_dims, + num_comps=values.shape[-1], + xlabel=xlabel, + ylabel=ylabel, + clabel=clabel) diff --git a/src/postgkyl/render/animate.py b/src/postgkyl/render/animate.py new file mode 100644 index 00000000..bd0622f0 --- /dev/null +++ b/src/postgkyl/render/animate.py @@ -0,0 +1,390 @@ +"""The canonical Matplotlib animation callable and its private helpers. + +``pg.animate``, ``operations.animate``, and the generated CLI all resolve to +the one public function in this module. It owns dataset grouping and +materialization as well as ``FuncAnimation`` / saved frames / movie compile. + +The module is separate from ``matplotlib.py`` because it owns the one +external-process dependency in this layer -- ``ffmpeg`` -- reached through +Matplotlib's ``FFMpegWriter``/``Animation.save``. Every entry point that needs +it resolves a binary via ``_ffmpeg.require_ffmpeg`` up front and raises a clear +``RuntimeError`` instead of failing deep inside the writer. +""" + +from __future__ import annotations + +import os.path +from collections.abc import Iterable +from typing import Annotated, TYPE_CHECKING + +import numpy as np + +from postgkyl.cli_spec import ( + CliType, + CommandSpec, + Execution, + PipelineInput, + ResultPolicy, + Section, + command, +) +from postgkyl.gdatastate import ( + GDataState, + group_blocks, + group_frames, + materialize_point_values, +) + +from . import matplotlib as backend +from ._ffmpeg import require_ffmpeg + +if TYPE_CHECKING: + from matplotlib.figure import Figure + +# Formats written through ffmpeg; PIL handles the rest (gif/webp/apng). +_VIDEO_EXTS = (".mp4", ".mov", ".avi", ".mkv") + + +def _normalize_frames(data, + *, + multiblock: bool = False) -> list[list["GDataState"]]: + """Group a flat input when requested and materialize every frame dataset.""" + items = list(data) + if items and all(isinstance(item, GDataState) for item in items): + items = group_frames(items) if multiblock else group_blocks(items) + frames = [([materialize_point_values(item)] if isinstance(item, GDataState) + else [materialize_point_values(dat) for dat in item]) + for item in items] + if not frames: + raise ValueError("animate: no datasets to animate.") + return frames + + +def _frame_value_range(frames: list[list["GDataState"]], + cutoff: float | None = None, + *, + yscale: float = 1.0, + zscale: float = 1.0) -> tuple[float, float]: + """Value range spanning every dataset in every frame. + + Each dataset is scaled by ``yscale`` (1-D) or ``zscale`` (2-D) before its + extrema are taken, matching the scale ``matplotlib.plot`` applies when it + actually draws the values -- otherwise a fixed range computed here would + not match the plotted (scaled) data. + + With ``cutoff`` (a central fraction in ``(0, 1]``), the range is clipped + to that percentile band of the per-dataset extrema instead of the true + min/max -- useful when a few outlier frames would otherwise wash out the + color/y-axis scale for the rest of the animation. + """ + extrema = [] + for frame in frames: + for dat in frame: + scaled = dat.values * (yscale if dat.num_dims == 1 else zscale) + extrema.append(np.nanmin(scaled)) + extrema.append(np.nanmax(scaled)) + extrema = np.array(extrema) + vmin, vmax = float(extrema.min()), float(extrema.max()) + if cutoff: + boundary = 100.0 * (1.0 - cutoff) / 2.0 + vmax = float(np.percentile(extrema, 100.0 - boundary)) + vmin = float(np.percentile(extrema, boundary)) + return vmin, vmax + + +def _draw_frame(frame: list["GDataState"], fig: "Figure", plot_kwargs: dict): + """Redraw one frame (a list of datasets drawn together) onto ``fig``. + + When the caller hasn't given an explicit ``title``, it is generated from + the first dataset's ``ctx`` (frame index and time) unless + ``plot_kwargs['notitle']`` is set; an explicit ``title`` is always + respected and shown on every frame. + """ + kwargs = dict(plot_kwargs) + notitle = kwargs.pop("notitle", False) + if not notitle and kwargs.get("title") is None: + dat0 = frame[0] + parts = [] + if dat0.ctx.get("frame") is not None: + parts.append(f"frame: {dat0.ctx['frame']:d}") + if dat0.ctx.get("time") is not None: + parts.append(f"time: {dat0.ctx['time']:.4e}") + kwargs["title"] = " ".join(parts) + return backend.plot(*frame, figure=fig, clear=True, no_show=True, **kwargs) + + +def _render_frame(index: int, frames: list[list["GDataState"]], fig: "Figure", + plot_kwargs: dict): + """``FuncAnimation``'s per-frame callback: draw ``frames[index]``.""" + return _draw_frame(frames[index], fig, plot_kwargs) + + +def _save_frame_worker(args) -> str: + """One frame, one process (see ``_save_frames``'s ``nproc`` path). Each + worker builds its own figure -- Matplotlib figures are not shared across + processes.""" + index, frame, plot_kwargs, prefix, dpi, figsize = args + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig = plt.figure(figsize=figsize) + try: + _draw_frame(frame, fig, plot_kwargs) + path = f"{prefix}_{index}.png" + fig.savefig(path, dpi=dpi) + finally: + plt.close(fig) + return path + + +def _save_frames(frames: list[list["GDataState"]], + prefix: str, + *, + dpi: int | None = None, + figsize=None, + plot_kwargs: dict | None = None, + nproc: int = 1) -> list[str]: + """Write ``_.png`` for every frame. + + Sequentially (``nproc == 1``), one figure is reused across every frame. + With ``nproc > 1``, frames are split across a :class:`multiprocessing.Pool` + of that many worker processes, each with its own figure. + """ + plot_kwargs = plot_kwargs or {} + if nproc > 1: + from multiprocessing import Pool + + args_list = [(i, frames[i], plot_kwargs, prefix, dpi, figsize) + for i in range(len(frames))] + with Pool(nproc) as pool: + return pool.map(_save_frame_worker, args_list) + + import matplotlib.pyplot as plt + + fig = plt.figure(figsize=figsize) + paths = [] + try: + for i in range(len(frames)): + _draw_frame(frames[i], fig, plot_kwargs) + path = f"{prefix}_{i}.png" + fig.savefig(path, dpi=dpi) + paths.append(path) + finally: + plt.close(fig) + return paths + + +def _compile_movie(frame_files: list[str], + output_file: str, + *, + fps: int | None = None, + duration: float = 100.0) -> None: + """Compile PNG frames into an animation: PIL for gif/webp/apng, the + Matplotlib ffmpeg writer for video containers. ``duration`` is the + per-frame time in milliseconds, used when ``fps`` is not given.""" + from PIL import Image + + ext = os.path.splitext(output_file)[1].lower() + if ext in (".gif", ".webp", ".apng"): + images = [Image.open(f) for f in frame_files] + images[0].save(output_file, + save_all=True, + append_images=images[1:], + duration=duration, + loop=0, + optimize=False) + return + if ext in _VIDEO_EXTS: + import matplotlib as mpl + import matplotlib.pyplot as plt + from matplotlib.animation import FFMpegWriter + + mpl.rcParams["animation.ffmpeg_path"] = require_ffmpeg("animate") + movie_fps = fps if fps else 1.0e3 / duration + writer = FFMpegWriter(fps=movie_fps) + with Image.open(frame_files[0]) as first: + width, height = first.size + dpi = 100 + fig = plt.figure(figsize=(width / dpi, height / dpi), dpi=dpi) + ax = fig.add_axes([0, 0, 1, 1]) + ax.axis("off") + try: + with writer.saving(fig, output_file, dpi): + for frame_file in frame_files: + ax.clear() + ax.axis("off") + with Image.open(frame_file) as frame: + ax.imshow(frame) + writer.grab_frame() + finally: + plt.close(fig) + return + raise ValueError(f"animate: unsupported output format {ext!r}") + + +@command( + CommandSpec(Section.RENDER, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) +def animate(data: Annotated[Iterable[GDataState | Iterable[GDataState]], + PipelineInput()], + *, + multiblock: bool = False, + grouptags: bool = False, + interval: int = 100, + variable_range: bool = False, + cutoffglobalrange: float | None = None, + notitle: bool = False, + no_show: bool = False, + save: bool = False, + saveas: str | None = None, + fps: int | None = None, + dpi: int | None = None, + saveframes: str | None = None, + figsize: Annotated[tuple[float, float] | str | None, + CliType(tuple[float, float] | None)] = None, + nproc: int = 1, + tmpdir: str | None = None): + """Animate a sequence of frames, one frame per dataset (or dataset group). + + Args: + data: a flat iterable of datasets (each becomes a single-dataset frame), + or an iterable of frames where each frame is itself a list of + datasets drawn together (overlaid, as in ``matplotlib.plot``). + multiblock: Force datasets with the same frame index into one frame. + grouptags: Build a separate animation for each dataset tag. + interval: live-animation delay between frames, in milliseconds. + variable_range: Recompute the value/color scale for every frame instead + of holding ``ymin``/``ymax``/``zmin``/``zmax`` constant. + cutoffglobalrange: clip the fixed range to this central percentile band + (see ``_frame_value_range``); ``None`` uses the true min/max. + notitle: suppress the per-frame frame/time title. + no_show: Do not open a live window (the ``FuncAnimation`` path only). + save: write to ``saveas`` (or ``anim.gif``) after building the frames. + saveas: output path; its extension selects the writer (``.gif``/ + ``.webp``/``.apng`` via PIL, ``.mp4``/``.mov``/``.avi``/``.mkv`` via + ffmpeg). + fps: frames per second for the saved movie; defaults from ``interval``. + dpi: resolution for saved frames/movies. + saveframes: when given, write ``_.png`` for every frame + instead of building a live ``FuncAnimation``. + figsize: figure size in inches, forwarded to ``matplotlib.plot``. + nproc: parallel worker processes for frame generation (``saveframes``, + or the ``tmpdir``-backed compile path below); ``1`` renders sequentially + in-process. + tmpdir: directory for the temporary frame directory used when ``nproc`` + is greater than 1 and ``saveframes`` is not given (frames are written + there, compiled into the output, then discarded). + + Returns: + The list of written frame paths when ``saveframes`` is set; otherwise + the ``FuncAnimation`` (keep a reference -- Matplotlib does not keep the + live animation alive for you). When ``nproc`` renders through the + ``tmpdir`` compile path, the compiled output path is returned instead. + + Raises: + ValueError: no datasets to animate, or an unsupported ``saveas`` + extension. + RuntimeError: saving to a video container without ffmpeg on ``PATH``. + """ + items = list(data) + if items and grouptags and all( + isinstance(item, GDataState) for item in items): + tags: dict[str, list[GDataState]] = {} + for item in items: + tags.setdefault(item.tag, []).append(item) + + def suffixed(path, tag): + if path is None: + return None + stem, extension = os.path.splitext(path) + return f"{stem}_{tag}{extension}" + + return [ + animate(tagged, + multiblock=multiblock, + interval=interval, + variable_range=variable_range, + cutoffglobalrange=cutoffglobalrange, + notitle=notitle, + no_show=no_show, + save=save, + saveas=suffixed(saveas, tag), + fps=fps, + dpi=dpi, + saveframes=suffixed(saveframes, tag), + figsize=figsize, + nproc=nproc, + tmpdir=tmpdir) for tag, tagged in tags.items() + ] + + frames = _normalize_frames(items, multiblock=multiblock) + plot_kwargs = {} + plot_kwargs["notitle"] = notitle + + if not variable_range: + vmin, vmax = _frame_value_range(frames, + cutoffglobalrange, + yscale=plot_kwargs.get("yscale", 1.0), + zscale=plot_kwargs.get("zscale", 1.0)) + # Applied as both the 1-D y-limits (ymin/ymax) and the 2-D color range + # (zmin/zmax) -- whichever the frame's dimensionality actually uses. + plot_kwargs.setdefault("ymin", vmin) + plot_kwargs.setdefault("ymax", vmax) + plot_kwargs.setdefault("zmin", vmin) + plot_kwargs.setdefault("zmax", vmax) + + num_frames = len(frames) + duration = 1.0e3 / fps if fps else float(interval) + out_file = saveas or "anim.gif" + if not os.path.splitext(out_file)[1]: + out_file += ".gif" + + if saveframes: + frame_files = _save_frames(frames, + saveframes, + dpi=dpi, + figsize=figsize, + plot_kwargs=plot_kwargs, + nproc=nproc) + if save or saveas: + _compile_movie(frame_files, out_file, fps=fps, duration=duration) + return frame_files + + if nproc > 1: + # No standing PNGs requested -- render into a scratch directory, compile, + # then discard it. Mirrors the ``saveframes`` path with parallel workers, + # so it always produces the compiled output (there is no live window to + # hand parallel workers' figures back to). + import tempfile + + with tempfile.TemporaryDirectory(dir=tmpdir) as tmp: + tmp_prefix = f"{tmp}/frame" + frame_files = _save_frames(frames, + tmp_prefix, + dpi=dpi, + figsize=figsize, + plot_kwargs=plot_kwargs, + nproc=nproc) + _compile_movie(frame_files, out_file, fps=fps, duration=duration) + return out_file + + import matplotlib.pyplot as plt + from matplotlib.animation import FuncAnimation + + fig = plt.figure(figsize=figsize) + anim = FuncAnimation(fig, + _render_frame, + num_frames, + fargs=(frames, fig, plot_kwargs), + interval=interval, + blit=False) + if save or saveas: + import matplotlib as mpl + + mpl.rcParams["animation.ffmpeg_path"] = require_ffmpeg("animate") + anim.save(out_file, writer="ffmpeg", fps=fps, dpi=dpi) + if not no_show: + plt.show() + return anim diff --git a/src/postgkyl/render/labels.py b/src/postgkyl/render/labels.py new file mode 100644 index 00000000..84ee4946 --- /dev/null +++ b/src/postgkyl/render/labels.py @@ -0,0 +1,87 @@ +"""LaTeX-ish label conversion for backends that cannot render mathtext. + +Matplotlib understands raw LaTeX-flavoured labels (``$z_0$``) natively via +mathtext, but Plotly and PyVista do not, so their labels are passed through +these converters instead. +""" + +from __future__ import annotations + +import re + +_LATEX_TO_UNICODE = { + r"\mu": "μ", + r"\nu": "ν", + r"\pi": "π", + r"\sigma": "σ", + r"\Sigma": "Σ", + r"\rho": "ρ", + r"\tau": "τ", + r"\chi": "χ", + r"\phi": "φ", + r"\psi": "ψ", + r"\omega": "ω", + r"\Omega": "Ω", + r"\alpha": "α", + r"\beta": "β", + r"\gamma": "γ", + r"\delta": "δ", + r"\Delta": "Δ", + r"\epsilon": "ε", + r"\zeta": "ζ", + r"\eta": "η", + r"\theta": "θ", + r"\Theta": "Θ", + r"\iota": "ι", + r"\kappa": "κ", + r"\lambda": "λ", + r"\Lambda": "Λ", + r"\parallel": "∥", + r"\perp": "⊥", +} + + +def latex_to_unicode(text: str) -> str: + """Convert common LaTeX commands (Greek letters, ``\\parallel``/``\\perp``) + to their Unicode characters, stripping a surrounding ``$...$``.""" + if not text: + return text + text = text.strip() + if text.startswith("$") and text.endswith("$"): + text = text[1:-1] + for latex, unicode_char in _LATEX_TO_UNICODE.items(): + text = text.replace(latex, unicode_char) + return text + + +def latex_to_html(text: str) -> str: + """Convert LaTeX subscripts and Greek letters to HTML. + + Plotly does not support LaTeX, but does support HTML, so this converts + common LaTeX syntax (``_{...}``/``_x``/Greek letters) to HTML equivalents. + """ + if not text: + return text + + text = text.strip() + if text.startswith("$") and text.endswith("$"): + text = text[1:-1] + + def _replace_latex_commands(value: str) -> str: + return latex_to_unicode(value) + + text = re.sub( + r'_\{([^{}]+)\}', + lambda match: f"{_replace_latex_commands(match.group(1))}", + text, + ) + text = re.sub( + r'_(\\[A-Za-z]+|[A-Za-z0-9])', + lambda match: f"{_replace_latex_commands(match.group(1))}", + text, + ) + text = _replace_latex_commands(text) + return text + + +__all__ = ["latex_to_html", "latex_to_unicode"] diff --git a/src/postgkyl/render/matplotlib.py b/src/postgkyl/render/matplotlib.py new file mode 100644 index 00000000..0a71fbd7 --- /dev/null +++ b/src/postgkyl/render/matplotlib.py @@ -0,0 +1,1312 @@ +"""The canonical Matplotlib plot function and its private drawing helpers. + +``pg.plot``, ``GData.plot``, ``operations.plot``, and the generated CLI are +aliases or lowerings of the one public function in this module. It owns every +plot option, dataset grouping, figure construction, saving, and display. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable +from contextlib import nullcontext +from typing import Annotated + +import matplotlib as mpl +import matplotlib.font_manager as fm +# Importing the toolkit registers Matplotlib's ``3d`` projection. +import mpl_toolkits.mplot3d # noqa: F401 +import numpy as np +from matplotlib import cm, colors, patches +from matplotlib.figure import Figure +from matplotlib.typing import ColorType + +from postgkyl.cli_spec import ( + CliType, + CommandSpec, + Execution, + KeyValue, + ResultPolicy, + Section, + command, +) +from postgkyl.gdatastate import ( + GDataState, + GDataStateGroup, + flatten_datasets, + group_blocks, + materialize_point_values, +) + +from ._prep import subplot_grid +from .style import apply_style + +_AXES_LABELS = [rf"$z_{i}$" for i in range(6)] +_INDEX_AXES_LABELS = [rf"$i_{i}$" for i in range(6)] +_OUTPUT_EXTENSIONS = (".png", ".pdf") +_AxisLimits = (tuple[float, float] | list[tuple[float, float]] + | dict[int, tuple[float, float]]) + + +def _indexed_saveas(saveas, index: int, indexed: bool): + """Add a family index to every requested output path when needed.""" + if saveas is None or not indexed: + return saveas + if isinstance(saveas, (str, os.PathLike)): + path = os.fspath(saveas) + stem, extension = os.path.splitext(path) + return f"{stem}_{index}{extension}" + return tuple(_indexed_saveas(path, index, True) for path in saveas) + + +def _default_output_stem(states) -> str: + """Best-effort output stem when ``save=True`` has no explicit path.""" + stems = [] + for i, data in enumerate(states): + file_name = getattr(data, "_file_name", "") or "" + if file_name: + stem = os.path.basename(file_name).split(".")[0] + else: + label = data.get_label() if hasattr(data, "get_label") else "" + stem = label.replace(" ", "_") if label else f"dataset_{i}" + stems.append(stem) + return "_".join(stems) or "matplotlib_output" + + +def _output_paths(save, saveas, states) -> tuple[str, ...]: + """Normalize and validate Matplotlib output paths. + + A sequence is accepted so the CLI can preserve combinations such as + ``--saveas plot.pdf --saveframes frame`` without saving outside the render + backend. Extension-less names retain the historical PNG default. + """ + empty_path = (isinstance(saveas, (str, os.PathLike)) + and not os.fspath(saveas)) + if saveas is None or empty_path: + if not save: + return () + paths = [_default_output_stem(states)] + elif isinstance(saveas, (str, os.PathLike)): + paths = [saveas] + else: + try: + paths = list(saveas) + except TypeError as err: + raise TypeError( + "'saveas' must be a path or an iterable of paths") from err + + normalized = [] + for path in paths: + try: + path = os.fspath(path) + except TypeError as err: + raise TypeError("every 'saveas' entry must be path-like") from err + _, ext = os.path.splitext(path) + ext = ext.lower() + if not ext: + path = f"{path}.png" + elif ext not in _OUTPUT_EXTENSIONS: + raise ValueError("Unsupported file format for saving. Supported formats " + "are: .png, .pdf") + normalized.append(path) + return tuple(normalized) + + +def _normalize_line_colors(color): + """Return ``color`` as a per-line tuple, or ``None`` for a scalar color. + + Matplotlib color specifications such as RGB/RGBA tuples are sequences too, + so test the complete value before interpreting it as a sequence of colors. + """ + if color is None or colors.is_color_like(color): + return None + if isinstance(color, str): + return None # let Matplotlib report its usual error for an invalid color + try: + line_colors = tuple(color) + except TypeError: + return None + if not line_colors: + raise ValueError("'color' must not be an empty sequence") + if not all(colors.is_color_like(line_color) for line_color in line_colors): + raise ValueError( + "every entry in a 'color' sequence must be a valid Matplotlib color") + return line_colors + + +def _normalize_linestyles(linestyle, num_datasets: int): + """Return one linestyle per dataset, or ``None`` for a scalar style.""" + if linestyle is None or isinstance(linestyle, str): + return None + # A Matplotlib custom dash pattern, e.g. ``(0, (5, 2))``, is one style + # despite being a sequence itself. + try: + is_dash_pattern = (len(linestyle) == 2 and np.isscalar(linestyle[0]) + and not isinstance(linestyle[1], str) + and all(np.isscalar(value) for value in linestyle[1])) + except (TypeError, IndexError): + is_dash_pattern = False + if is_dash_pattern: + return None + try: + linestyles = tuple(linestyle) + except TypeError: + return None + if not linestyles: + raise ValueError("'linestyle' must not be an empty sequence") + if len(linestyles) == 1: + return linestyles * num_datasets + if len(linestyles) != num_datasets: + raise ValueError( + f"'linestyle' contains {len(linestyles)} entries; expected either 1 " + f"(applied to every dataset) or {num_datasets} (one per dataset)") + return linestyles + + +def _pgkyl_colorbar(im, fig, ax, *, label: str = "", extend: str | None = None): + """The Postgkyl colorbar: appended beside ``ax`` (not shrinking it) via + ``make_axes_locatable``, instead of stealing width from the panel.""" + from mpl_toolkits.axes_grid1 import make_axes_locatable + + divider = make_axes_locatable(ax) + cax2 = divider.append_axes("right", size="3%", pad=0.05) + return fig.colorbar(im, cax=cax2, label=label or "", extend=extend) + + +def get_xkcd_safely(): + """An xkcd context manager + rc override that degrades gracefully when no + xkcd-style font is installed, instead of silently drawing with whatever + default font Matplotlib falls back to.""" + import warnings + + import matplotlib.pyplot as plt + + required_fonts = {"xkcd", "xkcd Script", "Comic Neue", "Comic Sans MS"} + available_fonts = {f.name for f in fm.fontManager.ttflist} + if required_fonts.isdisjoint(available_fonts): + warnings.warn( + "No xkcd-style font found (xkcd/xkcd Script/Comic Neue/Comic Sans " + "MS); falling back to the default sans-serif font.", + stacklevel=2) + font_rc = {"font.family": "sans-serif"} + else: + font_rc = {"font.family": "Comic Sans MS"} + return plt.xkcd, font_rc + + +def _nodal_grid(grid: list, cells: np.ndarray) -> list: + """Cell-center coordinates from nodal (edge) coordinate arrays. + + Handles both flat per-axis edge arrays and curvilinear (multi-dimensional, + ``.map()``-produced) coordinate arrays, where every coordinate array spans + all dimensions jointly. + """ + num_dims = len(grid) + if num_dims != len(cells): + raise ValueError("Number dimensions for 'grid' and 'values' doesn't match") + out = [] + for d in range(num_dims): + g = grid[d] + if g.ndim == 1: + if g.shape[0] == cells[d]: + out.append(g) + elif g.shape[0] == cells[d] + 1: + out.append(0.5 * (g[:-1] + g[1:])) + else: + raise ValueError("Something is terribly wrong...") + else: + if g.shape[d] == cells[d]: + out.append(g) + elif g.shape[d] == cells[d] + 1: + if num_dims == 1: + out.append(0.5 * (g[:-1] + g[1:])) + else: + out.append(0.5 * (g[:-1, :-1] + g[1:, 1:])) + else: + raise ValueError("Something is terribly wrong...") + return out + + +def _shared_component_range(states, zshift: float, zscale: float) -> list: + """Per-component ``(vmin, vmax)`` across *every* dataset drawn on one figure. + + When several 2-D datasets share one set of axes -- overlaid frames, or the + blocks of a multiblock field, each covering its own patch of the domain -- + each ``pcolormesh`` would otherwise normalize against only its own values, + so identical colors would mean different numbers in different patches and + the single shared colorbar would be a lie. Computing the range up front + makes one color scale describe the whole picture. + + Ranges are computed on the *plotted* values (``(v + zshift) * zscale``). + A component that is all-NaN (or absent from a dataset with fewer + components) yields ``(None, None)`` -- i.e. defer to Matplotlib. + """ + ranges = [] + num_comps = max(int(state.values.shape[-1]) for state in states) + for comp in range(num_comps): + low, high = np.inf, -np.inf + for state in states: + values = state.values + if comp >= values.shape[-1]: + continue + z = (values[..., comp] + zshift) * zscale + if not np.any(np.isfinite(z)): + continue + low = min(low, float(np.nanmin(z))) + high = max(high, float(np.nanmax(z))) + ranges.append((low, high) if np.isfinite(low) and low < high else (None, + None)) + return ranges + + +def _split_ylim_for_component(limits, comp: int): + """Resolve a split-axis y-limit specification for one component. + + ``limits`` may be one ``(min, max)`` pair shared by every component, a + sequence of pairs (one per component), or a mapping keyed by component + index. Missing sequence/mapping entries leave that component automatic. + """ + if limits is None: + return None + if isinstance(limits, dict): + limits = limits.get(comp) + if limits is None: + return None + else: + try: + is_shared_pair = (len(limits) == 2 + and all(value is None or np.isscalar(value) + for value in limits)) + except TypeError as err: + raise TypeError("split y-limits must be a (min, max) pair, a sequence " + "of pairs, or a component-indexed mapping") from err + if not is_shared_pair: + if comp >= len(limits): + return None + limits = limits[comp] + if limits is None: + return None + try: + if len(limits) != 2: + raise ValueError + except (TypeError, ValueError) as err: + raise ValueError("each split y-limit must be a (min, max) pair") from err + return tuple(limits) + + +def plot( + *datasets: GDataState | Iterable[GDataState], + multiblock: bool = False, + args: list[str] | None = None, + figure: Annotated[int | str | Figure | None, + CliType(int | None)] = None, + squeeze: bool = False, + transpose: bool = False, + grid_indices: bool = False, + num_axes: int | None = None, + start_axes: int = 0, + overlay_axes: bool = False, + num_subplot_row: int | None = None, + num_subplot_col: int | None = None, + streamline: bool = False, + sdensity: int = 1, + quiver: bool = False, + contour: bool = False, + clevels: str | None = None, + cnlevels: int | None = None, + cont_label: bool = False, + surface: bool = False, + comparison: bool = False, + alpha: float | None = None, + diverging: bool = False, + lineouts: int | None = None, + xmin: float | None = None, + xmax: float | None = None, + xscale: float = 1.0, + xshift: float = 0.0, + ymin: float | None = None, + ymax: float | None = None, + yscale: float = 1.0, + yshift: float = 0.0, + zmin: float | None = None, + zmax: float | None = None, + zscale: float = 1.0, + zshift: float = 0.0, + relax: bool = False, + style: str | None = None, + rcParams: Annotated[dict[str, object] | None, + CliType(dict[str, str] | None), + KeyValue()] = None, + no_legend: bool = False, + legend_labels: list[str] | None = None, + legend_subplot: int | None = None, + legend_loc: Annotated[str | int | tuple[float, float], + CliType(str)] = "best", + forcelegend: bool = False, + no_colorbar: bool = False, + xlabel: str | None = None, + ylabel: str | None = None, + clabel: str | None = None, + title: str | None = None, + subplot_titles: str | None = None, + subplot_xlabels: str | None = None, + subplot_ylabels: str | None = None, + logx: bool = False, + logy: bool = False, + logz: bool = False, + split_linear_log: bool = False, + split_point: float = 0.0, + split_log_side: str = "right", + split_width_ratios: tuple[float, float] = (1.0, 1.0), + split_gap: float = 0.0, + split_linear_ylim: Annotated[_AxisLimits | None, + CliType(tuple[float, float] | None)] = None, + split_log_ylim: Annotated[_AxisLimits | None, + CliType(tuple[float, float] | None)] = None, + no_split_right_ticks: bool = False, + split_legend_side: str = "log", + split_log_base: float = 10.0, + split_log_nonpositive: str = "clip", + split_seam_ticklabels: str = "left", + fixaspect: bool = False, + aspect: float | None = None, + edgecolors: str | None = None, + no_showgrid: bool = False, + hashtag: bool = False, + xkcd: bool = False, + color: Annotated[ColorType | Iterable[ColorType] | None, + CliType(str | None)] = None, + markersize: float | None = None, + linewidth: float | None = None, + linestyle: Annotated[str | Iterable[str] | None, + CliType(str | None)] = None, + figsize: Annotated[tuple[float, float] | str | None, + CliType(tuple[float, float] | None)] = None, + jet: bool = False, + cmap: str | None = None, + cval: float | None = None, + cval_min: float | None = None, + cval_max: float | None = None, + save: bool = False, + saveas: Annotated[str | os.PathLike | Iterable[str | os.PathLike] | None, + CliType(str | None)] = None, + dpi: int = 200, + no_show: bool = False, + clear: bool = False): + """Plot one or more datasets onto a shared figure and return it. + + Accepts ``plot(a)`` or ``plot(a, b, multiblock=True)``. The first dataset sets + the layout (dimensionality and panel count, after squeezing any size-1 axis + left by a coordinate ``select()``); every dataset (including the first) is + then drawn -- overlaid onto the same panels for 1-D, or onto the next + ``start_axes``-offset block of panels when ``num_axes`` spreads multiple + datasets' components across one grid (the old ``--subplots`` behaviour). + Set ``overlay_axes=True`` to keep every dataset in the *same* + ``start_axes`` block instead of advancing per dataset -- what the blocks of + one multiblock field want, since they are one field and belong in one panel. + + When more than one 2-D dataset is drawn as a ``pcolormesh`` and no explicit + ``zmin``/``zmax`` is given, all of them share one per-component color scale + (computed across the whole call) and one colorbar per panel, so the + colorbar describes every dataset on the axes rather than whichever was + drawn last. + + Most of the keyword arguments mirror main's ``output.plot``/CLI ``plot`` + 1:1 (contour/quiver/streamline/lineouts, shifts/scales, limits, labels, + legend, colorbar, aspect, log axes, xkcd/hashtag/jet, style). ``transpose`` + swaps the horizontal and vertical axes: in 1-D the coordinate moves to the + vertical axis; in 2-D the data, grid, and default labels are swapped before + drawing (shifts/scales keep their screen-axis meaning). ``grid_indices`` + replaces each plotted coordinate with its zero-based sample index without + changing the dataset itself. ``save``/``saveas``/ + ``no_show`` make the render call self-sufficient: ``saveas`` writes a PNG or + PDF according to its extension (an extension-less name defaults to PNG), + while ``save=True`` derives a PNG name from the input dataset. ``clear`` lets + ``render.animate`` redraw onto a persistent figure across frames. + A sequence of ``saveas`` paths writes the same figure to each path; this is + primarily useful to CLI callers that request both a named output and frame + output. + + For 1-D data, passing ``cmap`` together with ``cval`` colors the line by + mapping ``cval`` onto the colormap; ``cval_min``/``cval_max`` set the + normalization range (typically the min/max of the ``cval`` values across + all curves), so several curves drawn into the same axes share one scale. + ``color`` accepts either one Matplotlib color, applied to every line, or a + sequence containing one color per dataset (reused for all its components). + A sequence with one color per individual line is also accepted, in + dataset/component order. ``linestyle`` similarly accepts one style applied + to every dataset, or a sequence with one entry per dataset; a one-entry + sequence is broadcast to every dataset. + + For 2-D data, ``surface`` draws a 3D surface instead of a ``pcolormesh``. + When several 2-D datasets are overlaid onto the same axes for comparison, + set ``comparison`` so each surface/contour gets a distinct color and a + legend entry instead of overlapping and hiding each other. ``alpha`` + controls the surface transparency. Set ``legend_subplot`` to a zero-based + subplot index to draw the legend only there; ``legend_loc`` accepts any + Matplotlib legend location and defaults to ``"best"``. Explicit + ``legend_labels`` are used verbatim on every component, without an added + ``_cN`` suffix. ``xkcd`` no longer leaks into Matplotlib's global rcParams + past this call -- it is scoped to the figure drawn here. + + ``split_linear_log=True`` turns every 1-D component panel into a joined + pair split at ``split_point``: coordinates below the point are drawn on the + left and coordinates at/above it on the right. The right side is + logarithmic in y by default; ``split_log_side='left'`` reverses which half + is logarithmic. ``split_width_ratios`` and ``split_gap`` control the pair's + geometry. ``split_linear_ylim`` and ``split_log_ylim`` accept either one + ``(min, max)`` pair, a sequence of pairs (one per component), or a mapping + from component index to pair. ``split_legend_side`` is ``'linear'``, + ``'log'``, ``'left'``, or ``'right'``. This mode is intentionally limited + to 1-D, non-transposed plots. ``split_seam_ticklabels`` chooses which side + owns the label at the joined boundary: ``'left'`` (the default), + ``'right'``, ``'both'``, or ``'none'``. + + Args: + datasets: Datasets to draw. + multiblock: Force every dataset onto one figure instead of grouping fields. + args: Positional Matplotlib plot arguments. + figure: Existing or numbered figure to target. + squeeze: Remove singleton spatial axes before laying out panels. + transpose: Swap the horizontal and vertical display axes. + grid_indices: Plot zero-based sample indices instead of grid values. + num_axes: Number of logical component axes. + start_axes: First logical component axis to use. + overlay_axes: Draw every dataset in the same component-axis block. + num_subplot_row: Forced subplot row count. + num_subplot_col: Forced subplot column count. + streamline: Draw two-component fields as streamlines. + sdensity: Streamline density. + quiver: Draw two-component fields as arrows. + contour: Draw two-dimensional values as contours. + clevels: Explicit contour-level specification. + cnlevels: Number of contour levels. + cont_label: Label contour lines. + surface: Draw two-dimensional values as a three-dimensional surface. + comparison: Distinguish overlaid two-dimensional datasets. + alpha: Surface or comparison transparency. + diverging: Use a diverging colormap. + lineouts: Draw this many lineouts from two-dimensional data. + xmin: Lower horizontal-axis bound. + xmax: Upper horizontal-axis bound. + xscale: Horizontal-coordinate scale factor. + xshift: Horizontal-coordinate shift. + ymin: Lower vertical-axis bound. + ymax: Upper vertical-axis bound. + yscale: Vertical-coordinate scale factor. + yshift: Vertical-coordinate shift. + zmin: Lower value or color bound. + zmax: Upper value or color bound. + zscale: Value scale factor. + zshift: Value shift. + relax: Allow relaxed layout behavior for reused figures. + style: Matplotlib style name or style-file path. + rcParams: Matplotlib configuration overrides. + no_legend: Suppress legends for line plots. + legend_labels: Explicit dataset legend labels. + legend_subplot: Zero-based subplot receiving the legend. + legend_loc: Matplotlib legend location. + forcelegend: Draw a legend even for one unlabeled curve. + no_colorbar: Suppress color bars for field plots. + xlabel: Horizontal-axis label override. + ylabel: Vertical-axis label override. + clabel: Color-bar label override. + title: Figure-title override. + subplot_titles: Per-subplot title specification. + subplot_xlabels: Per-subplot horizontal-label specification. + subplot_ylabels: Per-subplot vertical-label specification. + logx: Use logarithmic horizontal coordinates. + logy: Use a logarithmic vertical axis. + logz: Use logarithmic values or colors. + split_linear_log: Split each one-dimensional panel into linear and log halves. + split_point: Coordinate joining the split halves. + split_log_side: Half using logarithmic scaling. + split_width_ratios: Relative widths of the split halves. + split_gap: Gap between split halves. + split_linear_ylim: Limits for the linear half. + split_log_ylim: Limits for the logarithmic half. + no_split_right_ticks: Suppress ticks on the right edge of split panels. + split_legend_side: Split half receiving the legend. + split_log_base: Logarithm base for the logarithmic half. + split_log_nonpositive: Handling of nonpositive logarithmic values. + split_seam_ticklabels: Half owning labels at the split seam. + fixaspect: Use equal physical scaling on coordinate axes. + aspect: Explicit axes aspect ratio. + edgecolors: Mesh edge color. + no_showgrid: Suppress plot grid lines. + hashtag: Prefix labels with a hash marker. + xkcd: Draw using Matplotlib's XKCD context. + color: Line color or per-line colors. + markersize: Line-marker size. + linewidth: Line width. + linestyle: Line style or per-dataset styles. + figsize: Figure width and height in inches. + jet: Use the legacy jet colormap. + cmap: Matplotlib colormap name. + cval: Scalar used to color a one-dimensional curve. + cval_min: Lower normalization bound for curve coloring. + cval_max: Upper normalization bound for curve coloring. + save: Save to an automatically derived output name. + saveas: Explicit image output path or paths. + dpi: Saved-image resolution. + no_show: Do not display the figures interactively. + clear: Clear a reused figure before drawing. + + Returns: + One Matplotlib ``Figure``, or one figure per distinct field/frame family. + + Raises: + ValueError: nothing to plot, a dataset has no values, a dataset has more + than two (squeezed) dimensions, or (without ``squeeze``) a reused + figure does not have enough axes for the panel count. + """ + import matplotlib.pyplot as plt + + states = flatten_datasets(datasets) + if not states: + raise ValueError("nothing to plot") + group_call = len(datasets) == 1 and isinstance(datasets[0], GDataStateGroup) + families = ([states] if multiblock or group_call or figure is not None else + group_blocks(states)) + indexed = len(families) > 1 or (save and saveas is not None and isinstance( + saveas, + (str, os.PathLike)) and not os.path.splitext(os.fspath(saveas))[1]) + figures = [] + plot_args = () if args is None else args + for family_index, states in enumerate(families): + states = [materialize_point_values(state) for state in states] + family_saveas = _indexed_saveas(saveas, family_index, indexed) + for st in states: + if st.values is None: + raise ValueError("dataset has no values to plot") + + line_colors = _normalize_line_colors(color) + line_styles = _normalize_linestyles(linestyle, len(states)) + + # ---- Style / global rcParams novelties ---- + apply_style(style) if style else apply_style("postgkyl") + if rcParams: + for key, value in rcParams.items(): + mpl.rcParams[key] = value + if cmap: + mpl.rcParams["image.cmap"] = cmap + elif diverging: + mpl.rcParams["image.cmap"] = "RdBu_r" + if jet: # not for general use -- only for comparing against literature + mpl.rcParams["image.cmap"] = "jet" + if xkcd: + xkcd_cm, xkcd_rc = get_xkcd_safely() + else: + xkcd_cm, xkcd_rc = nullcontext, {} + if color is not None and line_colors is None: + mpl.rcParams["lines.color"] = color + if linewidth: + mpl.rcParams["lines.linewidth"] = linewidth + if linestyle is not None and line_styles is None: + mpl.rcParams["lines.linestyle"] = linestyle + + with xkcd_cm(), mpl.rc_context(rc=xkcd_rc): + + if not aspect: + aspect = 1.0 + + # ---- Phase 1: figure/axes layout, from the first dataset ---- + ref = states[0] + ref_cells = ref.num_cells + ref_num_dims = len(ref_cells) - int(np.sum(ref_cells <= 1)) + if ref_num_dims > 2: + raise ValueError("Only 1D and 2D plots are currently supported") + if line_colors is not None and ref_num_dims != 1: + raise ValueError("a 'color' sequence is only supported for 1D plots") + line_colors_by_dataset = False + if line_colors is not None: + component_step = 2 if (streamline or quiver) else 1 + expected_colors = sum(st.values.shape[-1] // component_step + for st in states) + if len(line_colors) == len(states): + line_colors_by_dataset = True + elif len(line_colors) != expected_colors: + raise ValueError( + f"'color' contains {len(line_colors)} entries; expected either " + f"{len(states)} (one per dataset) or {expected_colors} " + "(one per line)") + if split_linear_log: + if ref_num_dims != 1: + raise ValueError("'split_linear_log' is only supported for 1D plots") + if transpose: + raise ValueError( + "'split_linear_log' cannot be combined with 'transpose'") + if logy: + raise ValueError("'logy' is redundant with 'split_linear_log'; use " + "'split_log_side' to choose the logarithmic half") + if split_log_side not in {"left", "right"}: + raise ValueError("'split_log_side' must be 'left' or 'right'") + if split_legend_side not in {"linear", "log", "left", "right"}: + raise ValueError("'split_legend_side' must be 'linear', 'log', " + "'left', or 'right'") + if split_log_nonpositive not in {"clip", "mask"}: + raise ValueError("'split_log_nonpositive' must be 'clip' or 'mask'") + if split_seam_ticklabels not in {"left", "right", "both", "none"}: + raise ValueError("'split_seam_ticklabels' must be 'left', 'right', " + "'both', or 'none'") + try: + split_point = float(split_point) + except (TypeError, ValueError) as err: + raise TypeError("'split_point' must be a finite number") from err + if not np.isfinite(split_point): + raise ValueError("'split_point' must be a finite number") + try: + split_log_base = float(split_log_base) + except (TypeError, ValueError) as err: + raise TypeError("'split_log_base' must be a number") from err + if not np.isfinite( + split_log_base) or split_log_base <= 0 or split_log_base == 1: + raise ValueError( + "'split_log_base' must be positive and not equal to 1") + try: + split_width_ratios = tuple(float(v) for v in split_width_ratios) + if (len(split_width_ratios) != 2 + or any(not np.isfinite(v) or v <= 0 for v in split_width_ratios)): + raise ValueError + except (TypeError, ValueError) as err: + raise ValueError( + "'split_width_ratios' must contain two positive values") from err + try: + split_gap = float(split_gap) + except (TypeError, ValueError) as err: + raise TypeError("'split_gap' must be a number") from err + if not np.isfinite(split_gap) or split_gap < 0: + raise ValueError("'split_gap' must be non-negative") + + # Surface plots need 3D axes; only meaningful for 2D data. + use_3d = bool(surface) and ref_num_dims == 2 + subplot_kw = {"projection": "3d"} if use_3d else {} + + coordinate_labels = _INDEX_AXES_LABELS if grid_indices else _AXES_LABELS + default_xlabel, default_ylabel = coordinate_labels[0], coordinate_labels[ + 1] + if transpose and ref_num_dims == 2: + # The data axes are swapped before drawing, so the default label base + # names swap too; the shift/scale annotations below keep their + # screen-axis meaning (xshift still shifts the horizontal axis). + default_xlabel, default_ylabel = default_ylabel, default_xlabel + layout_xlabel = xlabel + layout_ylabel = ylabel + layout_clabel = clabel + if layout_xlabel is None: + layout_xlabel = default_xlabel if lineouts != 1 else coordinate_labels[1] + if xshift != 0.0 and xscale != 1.0: + layout_xlabel = rf"({layout_xlabel:s} + {xshift:.2e}) $\times$ {xscale:.2e}" + elif xshift != 0.0: + layout_xlabel = rf"{layout_xlabel:s} + {xshift:.2e}" + elif xscale != 1.0: + layout_xlabel = rf"{layout_xlabel:s} $\times$ {xscale:.2e}" + if layout_ylabel is None and ref_num_dims == 2 and lineouts is None: + layout_ylabel = default_ylabel + # NB: these elif conditions check xshift/xscale, not yshift/yscale -- + # a literal main bug (commands.plot's ylabel branch), kept for fidelity. + if yshift != 0.0 and yscale != 1.0: + layout_ylabel = rf"({layout_ylabel:s} + {yshift:.2e}) $\times$ {yscale:.2e}" + elif xshift != 0.0: + layout_ylabel = rf"{layout_ylabel:s} + {yshift:.2e}" + elif xscale != 1.0: + layout_ylabel = rf"{layout_ylabel:s} $\times$ {yscale:.2e}" + if zscale != 1.0: + layout_clabel = (rf"{layout_clabel:s} $\times$ {zscale:.3e}" + if layout_clabel else rf"$\times$ {zscale:.3e}") + if transpose and ref_num_dims == 1: + # The coordinate moves to the vertical axis, so the (resolved) labels + # follow it -- including the shift/scale annotation, which travels with + # the data it describes. + layout_xlabel, layout_ylabel = layout_ylabel, layout_xlabel + + if isinstance(figsize, str): + parts = figsize.split(",") + figsize = (float(parts[0]), float(parts[1])) + + if figure is None: + mpl_fig = plt.figure(figsize=figsize) + elif isinstance(figure, int): + mpl_fig = plt.figure(figure, figsize=figsize) + elif isinstance(figure, mpl.figure.Figure): + mpl_fig = figure + elif isinstance(figure, str): + mpl_fig = plt.figure(int(figure), figsize=figsize) + else: + raise TypeError( + "'figure' keyword needs to be one of None (default), int, str, " + "or a Matplotlib Figure") + if clear: + mpl_fig.clf() + + step = 2 if (streamline or quiver) else 1 + ref_idx_comps = range(int(np.floor(ref.num_comps / step))) + layout_num_comps = num_axes if num_axes else len(ref_idx_comps) + + physical_num_axes = (2 * (1 if squeeze else layout_num_comps) + if split_linear_log else + (1 if squeeze else layout_num_comps)) + if mpl_fig.axes: + ax = mpl_fig.axes + if physical_num_axes > len(ax): + raise ValueError("Trying to plot into figure with not enough axes") + else: + if split_linear_log: + # Each logical component owns a nested 1x2 GridSpec. Nesting keeps + # the within-pair gap independent from spacing between components. + logical_num_axes = 1 if squeeze else layout_num_comps + num_rows, num_cols = ((1, 1) if squeeze else subplot_grid( + logical_num_axes, num_subplot_row, num_subplot_col)) + outer = mpl_fig.add_gridspec(num_rows, num_cols) + ax = [] + shared_left = None + shared_right = None + for logical_idx in range(logical_num_axes): + row, col = divmod(logical_idx, num_cols) + inner = outer[row, col].subgridspec(1, + 2, + width_ratios=split_width_ratios, + wspace=split_gap) + left_ax = mpl_fig.add_subplot(inner[0], sharex=shared_left) + right_ax = mpl_fig.add_subplot(inner[1], sharex=shared_right) + if shared_left is None: + shared_left, shared_right = left_ax, right_ax + ax.extend((left_ax, right_ax)) + + if title: + mpl_fig.suptitle(title) + if layout_xlabel: + mpl_fig.supxlabel(layout_xlabel) + if layout_ylabel: + mpl_fig.supylabel(layout_ylabel) + sub_titles = subplot_titles.split(",") if subplot_titles else [] + sub_xlabels = subplot_xlabels.split(",") if subplot_xlabels else [] + sub_ylabels = subplot_ylabels.split(",") if subplot_ylabels else [] + pair_center = (split_width_ratios[0] + + split_width_ratios[1]) / (2.0 * split_width_ratios[0]) + for logical_idx in range(logical_num_axes): + left_ax, right_ax = ax[2 * logical_idx:2 * logical_idx + 2] + sub_title = (sub_titles[logical_idx] + if logical_idx < len(sub_titles) else "") + sub_xlabel = (sub_xlabels[logical_idx] + if logical_idx < len(sub_xlabels) else "") + sub_ylabel = (sub_ylabels[logical_idx] + if logical_idx < len(sub_ylabels) else "") + left_ax.set_ylabel(sub_ylabel) + if sub_xlabel: + left_ax.set_xlabel(sub_xlabel) + left_ax.xaxis.set_label_coords(pair_center, -0.1) + if sub_title: + left_ax.set_title(sub_title, x=pair_center, y=1.08) + if not no_split_right_ticks: + right_ax.yaxis.tick_right() + right_ax.yaxis.set_label_position("right") + elif squeeze: # Plotting into 1 panel + mpl_fig.subplots(1, 1, subplot_kw=subplot_kw) + ax = mpl_fig.axes + ax[0].set_xlabel(layout_xlabel) + ax[0].set_ylabel(layout_ylabel) + if title is not None: + ax[0].set_title(title, y=1.08) + else: # Plotting each component into its own subplot + num_rows, num_cols = subplot_grid(layout_num_comps, num_subplot_row, + num_subplot_col) + if ref_num_dims == 1 or lineouts is not None: + mpl_fig.subplots(num_rows, + num_cols, + sharex=True, + subplot_kw=subplot_kw) + elif use_3d: # 3D axes cannot share x/y with each other + mpl_fig.subplots(num_rows, num_cols, subplot_kw=subplot_kw) + else: # In 2D, share y-axis as well + mpl_fig.subplots(num_rows, num_cols, sharex=True, sharey=True) + ax = mpl_fig.axes + for extra in ax[layout_num_comps:]: + extra.axis("off") + if title: + mpl_fig.suptitle(title) + if layout_xlabel: + mpl_fig.supxlabel(layout_xlabel) + if layout_ylabel: + mpl_fig.supylabel(layout_ylabel) + + for ax_idx in range(len(ax)): + sub_titles = subplot_titles.split(",") if subplot_titles else [] + sub_xlabels = subplot_xlabels.split(",") if subplot_xlabels else [] + sub_ylabels = subplot_ylabels.split(",") if subplot_ylabels else [] + sub_title = sub_titles[ax_idx] if ax_idx < len(sub_titles) else "" + sub_xlabel = sub_xlabels[ax_idx] if ax_idx < len( + sub_xlabels) else "" + sub_ylabel = sub_ylabels[ax_idx] if ax_idx < len( + sub_ylabels) else "" + ax[ax_idx].set_xlabel(sub_xlabel) + ax[ax_idx].set_ylabel(sub_ylabel) + if sub_title: + ax[ax_idx].set_title(sub_title, y=1.08) + + # One color scale for every dataset drawn here (see + # _shared_component_range). Only the plain 2-D pcolormesh path consumes + # it: surface/contour/quiver/streamline/lineouts each own their own + # normalization, and an explicit zmin/zmax always wins. + shared_z = None + if (len(states) > 1 and ref_num_dims == 2 and zmin is None + and zmax is None + and not (surface or contour or quiver or streamline or diverging) + and lineouts is None): + shared_z = _shared_component_range(states, zshift, zscale) + + if legend_subplot is not None: + num_legend_subplots = 1 if squeeze else layout_num_comps + if not isinstance(legend_subplot, int): + raise TypeError("'legend_subplot' must be an integer or None") + if legend_subplot < 0 or legend_subplot >= num_legend_subplots: + raise ValueError( + f"'legend_subplot' must be between 0 and {num_legend_subplots - 1}" + ) + + # ---- Phase 2: draw each dataset ---- + im = None + cur_start_axes = start_axes + line_color_idx = 0 + for ds_i, data in enumerate(states): + if legend_labels is not None and ds_i < len(legend_labels): + label_prefix = legend_labels[ds_i] + explicit_legend_label = True + elif len(states) > 1 or forcelegend: + label_prefix = data.get_label() + explicit_legend_label = False + else: + label_prefix = "" + explicit_legend_label = False + + cells = data.num_cells + grid = list(data.grid) + values = data.values + num_dims = len(cells) - int(np.sum(cells <= 1)) + if num_dims > 2: + raise ValueError("Only 1D and 2D plots are currently supported") + if split_linear_log and num_dims != 1: + raise ValueError( + "every dataset must be 1D when 'split_linear_log' is set") + + axes_labels = list(coordinate_labels) + if len(grid) > num_dims: + idx = [d for d in range(len(grid)) if cells[d] <= 1] + grid = [g.squeeze() for g in grid] + if idx: + for d in reversed(idx): + grid.pop(d) + cells = np.delete(cells, idx) + axes_labels = list(np.delete(np.array(axes_labels), idx)) + values = np.squeeze(values, tuple(idx)) + if grid and grid[0].ndim > 1: # curvilinear (mapped) coordinates + for d in range(num_dims): + for i in reversed(idx): + grid[d] = np.mean(grid[d], axis=i) + + if transpose and num_dims == 2: # swap the horizontal and vertical axes + values = np.swapaxes(values, 0, 1) + g0, g1 = grid[1], grid[0] + if g0.ndim > 1: # curvilinear coordinate arrays span both axes jointly + g0, g1 = g0.transpose(), g1.transpose() + grid[0], grid[1] = g0, g1 + cells = cells[[1, + 0]] # fancy indexing: num_cells may alias ctx["cells"] + axes_labels[0], axes_labels[1] = axes_labels[1], axes_labels[0] + + if grid_indices: + grid = [np.arange(int(num_cells)) for num_cells in cells] + + num_comps = values.shape[-1] + idx_comps = range(int(np.floor(num_comps / step))) + + for comp in idx_comps: + logical_ax_idx = 0 if squeeze else comp + cur_start_axes + if split_linear_log: + component_axes = ax[2 * logical_ax_idx:2 * logical_ax_idx + 2] + cax = component_axes[0] + else: + cax = ax[logical_ax_idx] + component_axes = [cax] + comp_label = (label_prefix if explicit_legend_label else + (f"{label_prefix:s}_c{comp:d}".strip("_") + if len(idx_comps) > 1 else label_prefix)) + comp_legend = (not no_legend and + (legend_subplot is None or + (logical_ax_idx == legend_subplot + if split_linear_log else cax is ax[legend_subplot]))) + comp_colorbar = not no_colorbar + + if num_dims == 1: + nodal_grid = _nodal_grid(grid, cells) + x = (nodal_grid[0] + xshift) * xscale + y = (values[..., comp] + yshift) * yscale + if transpose: # put the coordinate on the vertical axis + x, y = y, x + # Color the line from the colormap when a 'cval' is given (1D only). + if line_colors is None: + line_color = color + elif line_colors_by_dataset: + line_color = line_colors[ds_i] + else: + line_color = line_colors[line_color_idx] + line_color_idx += 1 + if cmap and cval is not None: + if cval_max is not None and cval_min is not None and cval_max != cval_min: + t = (cval - cval_min) / (cval_max - cval_min) + else: + t = 0.5 + line_color = plt.get_cmap(cmap)(t) + line_style = line_styles[ds_i] if line_styles is not None else None + line_kwargs = dict(color=line_color, + label=comp_label, + markersize=markersize) + if line_style is not None: + line_kwargs["linestyle"] = line_style + if split_linear_log: + left_mask = x < split_point + split_masks = (left_mask, ~left_mask) + im = [] + for split_ax, mask in zip(component_axes, split_masks): + im.extend( + split_ax.plot(x[mask], y[mask], *plot_args, **line_kwargs)) + else: + im = cax.plot(x, y, *plot_args, **line_kwargs) + # Add a colorbar describing the cval-to-color mapping once per axes. + if (cmap and cval is not None and comp_colorbar + and cval_max is not None and cval_min is not None + and cval_max != cval_min + and not getattr(cax, "_pgkyl_cval_cbar", False)): + mappable = cm.ScalarMappable(norm=colors.Normalize(vmin=cval_min, + vmax=cval_max), + cmap=plt.get_cmap(cmap)) + _pgkyl_colorbar(mappable, mpl_fig, cax, label=layout_clabel) + cax._pgkyl_cval_cbar = True + + elif num_dims == 2: + extend = None + + if surface: # ------------------------------------------------------ + nodal_grid = _nodal_grid(grid, cells) + xg = (nodal_grid[0] + xshift) * xscale + yg = (nodal_grid[1] + yshift) * yscale + z = (values[..., comp].transpose() + zshift) * zscale + if xg.ndim == 1: + xg, yg = np.meshgrid(xg, yg) + else: + xg, yg = xg.transpose(), yg.transpose() + # Count how many overlays already live on these axes so each gets + # a distinct color (used for both surface and contour comparisons). + overlay_count = getattr(cax, "_pgkyl_overlay_count", 0) + cax._pgkyl_overlay_count = overlay_count + 1 + if comparison or bool(color): + surf_color = color if bool(color) else f"C{overlay_count:d}" + im = cax.plot_surface(xg, + yg, + z, + color=surf_color, + alpha=alpha if alpha is not None else 0.6, + linewidth=0, + antialiased=True, + shade=True) + if comp_label: + handles = getattr(cax, "_pgkyl_handles", []) + handles.append( + patches.Patch(color=surf_color, label=comp_label)) + cax._pgkyl_handles = handles + else: + im = cax.plot_surface(xg, + yg, + z, + cmap=mpl.rcParams["image.cmap"], + alpha=alpha if alpha is not None else 1.0, + linewidth=0, + antialiased=True) + if comp_colorbar: + mpl_fig.colorbar(im, + ax=cax, + label=layout_clabel or "", + shrink=0.6, + pad=0.1) + if layout_clabel: + cax.set_zlabel(layout_clabel) + if zmin is not None or zmax is not None: + cax.set_zlim(zmin, zmax) + comp_colorbar = False + + elif contour: # ------------------------------------------------------ + levels = 10 + if cnlevels: + levels = int(cnlevels) - 1 + elif clevels: + if ":" in clevels: + s = clevels.split(":") + levels = np.linspace(float(s[0]), float(s[1]), int(s[2])) + else: + levels = np.array(clevels.split(",")) + levels = np.array(list(filter(None, levels))) + if isinstance(levels, np.ndarray) and len(levels) == 1: + comp_colorbar = False + nodal_grid = _nodal_grid(grid, cells) + x = (nodal_grid[0] + xshift) * xscale + y = (nodal_grid[1] + yshift) * yscale + z = (values[..., comp].transpose() + zshift) * zscale + cont_colors = color + if comparison and not bool(color): + # Give each overlaid dataset a distinct, single color + legend entry. + overlay_count = getattr(cax, "_pgkyl_overlay_count", 0) + cax._pgkyl_overlay_count = overlay_count + 1 + cont_colors = f"C{overlay_count:d}" + if comp_label: + handles = getattr(cax, "_pgkyl_handles", []) + handles.append( + patches.Patch(color=cont_colors, label=comp_label)) + cax._pgkyl_handles = handles + comp_colorbar = False + im = cax.contour(x, + y, + z, + levels, + *plot_args, + origin="lower", + colors=cont_colors, + linewidths=linewidth) + if cont_label: + cax.clabel(im, inline=1) + + elif quiver: # ----------------------------------------------------- + skip = int(np.max((len(grid[0]), len(grid[1]))) // 15) + skip2 = int(skip // 2) + nodal_grid = _nodal_grid(grid, cells) + if nodal_grid[0].ndim == 1: + x = (nodal_grid[0][skip2::skip] + xshift) * xscale + y = (nodal_grid[1][skip2::skip] + yshift) * yscale + else: + x = (nodal_grid[0][skip2::skip, skip2::skip] + xshift) * xscale + y = (nodal_grid[1][skip2::skip, skip2::skip] + yshift) * yscale + z1 = (values[skip2::skip, skip2::skip, 2 * comp].transpose() + + zshift) * zscale + z2 = (values[skip2::skip, skip2::skip, 2 * comp + 1].transpose() + + zshift) * zscale + im = cax.quiver(x, y, z1, z2) + + elif streamline: # ------------------------------------------------- + if color: + cl = color + else: + cl = np.sqrt(values[..., 2 * comp]**2 + + values[..., 2 * comp + 1]**2).transpose() + nodal_grid = _nodal_grid(grid, cells) + x = (nodal_grid[0] + xshift) * xscale + y = (nodal_grid[1] + yshift) * yscale + z1 = (values[..., 2 * comp].transpose() + zshift) * zscale + z2 = (values[..., 2 * comp + 1].transpose() + zshift) * zscale + im = cax.streamplot(x, + y, + z1, + z2, + *plot_args, + density=sdensity, + broken_streamlines=False, + color=cl, + linewidth=linewidth) + + elif lineouts is not None: # --------------------------------------- + num_lines = values.shape[1] if lineouts == 0 else values.shape[0] + nodal_grid = _nodal_grid(grid, cells) + + if lineouts == 0: + x = (nodal_grid[0] + xshift) * xscale + line_vmin = (nodal_grid[1][0] + yshift) * yscale + line_vmax = (nodal_grid[1][-1] + yshift) * yscale + cbar_label = clabel or axes_labels[1] + else: + x = (nodal_grid[1] + xshift) * xscale + line_vmin = (nodal_grid[0][0] + yshift) * yscale + line_vmax = (nodal_grid[0][-1] + yshift) * yscale + cbar_label = clabel or axes_labels[0] + line_idx = [slice(0, u) for u in values.shape] + line_idx[-1] = comp + for line in range(num_lines): + line_color = cm.inferno(line / (num_lines - 1)) + if lineouts == 0: + line_idx[1] = line + else: + line_idx[0] = line + y = (values[tuple(line_idx)] + yshift) * yscale + im = cax.plot(x, y, *plot_args, color=line_color) + mappable = cm.ScalarMappable(norm=colors.Normalize(vmin=line_vmin, + vmax=line_vmax, + clip=False), + cmap=cm.inferno) + _pgkyl_colorbar(mappable, mpl_fig, cax, label=cbar_label) + comp_colorbar = False + comp_legend = False + + else: # ------------------------------------------------------------ + if zmin is not None and zmax is not None: + extend = "both" + elif zmax is not None: + extend = "max" + elif zmin is not None: + extend = "min" + x = (grid[0] + xshift) * xscale + y = (grid[1] + yshift) * yscale + z = (values[..., comp].transpose() + zshift) * zscale + if len(x) == z.shape[1] or len(y) == z.shape[0]: + nodal_grid = _nodal_grid(grid, cells) + x = (nodal_grid[0] + xshift) * xscale + y = (nodal_grid[1] + yshift) * yscale + if x.ndim > 1: + x, y = x.transpose(), y.transpose() + comp_zmin, comp_zmax = zmin, zmax + if diverging: + comp_zmax = np.abs(z).max() + comp_zmin = -comp_zmax + elif shared_z is not None and comp < len(shared_z): + comp_zmin, comp_zmax = shared_z[comp] + vmax, vmin = comp_zmax, comp_zmin + norm = None + if logz: + if diverging: + tmp = vmax / 1000 + norm = colors.SymLogNorm(linthresh=tmp, + linscale=tmp, + vmin=vmin, + vmax=vmax, + base=10) + else: + norm = colors.LogNorm(vmin=vmin, vmax=vmax) + vmin, vmax = None, None + im = cax.pcolormesh(x, + y, + z, + norm=norm, + vmin=vmin, + vmax=vmax, + edgecolors=edgecolors, + linewidth=0.1, + shading="auto", + *plot_args) + # One colorbar per (panel, component), not one per dataset drawn + # into it: with several datasets on shared axes (multiblock blocks, + # overlays) the per-dataset call used to stack an identical + # colorbar per dataset, each shrinking the figure further. They + # share one scale now, so the first describes them all. Keying on + # the component too keeps ``squeeze``'s several components in one + # panel getting their own (genuinely differently scaled) colorbars. + drawn = getattr(cax, "_pgkyl_cbar_comps", None) + if drawn is None: + drawn = cax._pgkyl_cbar_comps = set() + if not color and comp_colorbar and not streamline and comp not in drawn: + _pgkyl_colorbar(im, + mpl_fig, + cax, + extend=extend, + label=layout_clabel) + drawn.add(comp) + else: + raise ValueError(f"{num_dims:d}D data not supported") + + legend_ax = cax + if split_linear_log: + if split_legend_side == "left": + legend_ax = component_axes[0] + elif split_legend_side == "right": + legend_ax = component_axes[1] + elif split_legend_side == "linear": + legend_ax = component_axes[0 if split_log_side == "right" else 1] + else: # log + legend_ax = component_axes[0 if split_log_side == "left" else 1] + if comp_legend: + if getattr(legend_ax, "_pgkyl_handles", None): + # Overlaid 2D datasets (surface/contour comparison): real legend. + legend_ax.legend(handles=legend_ax._pgkyl_handles, loc=legend_loc) + elif num_dims == 1 and comp_label != "": + legend_ax.legend(loc=legend_loc) + elif not (surface and num_dims == 2): + legend_ax.text(0.03, + 0.96, + comp_label, + bbox={ + "facecolor": "w", + "edgecolor": "w", + "alpha": 0.8, + "boxstyle": "round" + }, + verticalalignment="top", + horizontalalignment="left", + transform=legend_ax.transAxes) + for side_idx, side_ax in enumerate(component_axes): + side_ax.grid(not no_showgrid) + if hashtag and (not split_linear_log or side_ax is legend_ax): + side_ax.text(0.97, + 0.03, + "#pgkyl", + bbox={ + "facecolor": "w", + "edgecolor": "w", + "alpha": 0.8, + "boxstyle": "round" + }, + verticalalignment="bottom", + horizontalalignment="right", + transform=side_ax.transAxes) + if logx: + side_ax.set_xscale("log") + if logy: + side_ax.set_yscale("log") + if split_linear_log: + is_log_side = ((side_idx == 0 and split_log_side == "left") + or (side_idx == 1 and split_log_side == "right")) + if is_log_side: + side_ax.set_yscale("log", + base=split_log_base, + nonpositive=split_log_nonpositive) + side_ylim = _split_ylim_for_component(split_log_ylim, comp) + else: + side_ax.set_yscale("linear") + side_ylim = _split_ylim_for_component(split_linear_ylim, comp) + else: + side_ylim = None + if num_dims == 1 and not relax: # this causes troubles with contours + side_ax.autoscale(enable=True, axis="x", tight=True) + side_ax.autoscale(enable=True, axis="y") + if split_linear_log: + if side_idx == 0: + side_ax.set_xlim(xmin, split_point) + else: + side_ax.set_xlim(split_point, xmax) + if not logx: + prune = None + if ((split_seam_ticklabels == "left" and side_idx == 1) + or (split_seam_ticklabels == "right" and side_idx == 0) + or split_seam_ticklabels == "none"): + prune = "upper" if side_idx == 0 else "lower" + if prune is not None: + side_ax.xaxis.get_major_locator().set_params(prune=prune) + elif xmin is not None or xmax is not None: + side_ax.set_xlim(xmin, xmax) + if ymin is not None or ymax is not None: + side_ax.set_ylim(ymin, ymax) + if side_ylim is not None: + side_ax.set_ylim(*side_ylim) + if fixaspect and not (surface and num_dims == 2): + plt.setp(side_ax, aspect=aspect) + + if num_axes and not overlay_axes: + cur_start_axes += num_comps + + mpl_fig.tight_layout() + for output_path in _output_paths(save, family_saveas, states): + mpl_fig.savefig(output_path, dpi=dpi) + figures.append(mpl_fig) + if not no_show: + plt.show() + return figures[0] if len(figures) == 1 else figures + + +command( + CommandSpec(Section.RENDER, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT))(plot) diff --git a/src/postgkyl/render/plotly.py b/src/postgkyl/render/plotly.py new file mode 100644 index 00000000..81b114dc --- /dev/null +++ b/src/postgkyl/render/plotly.py @@ -0,0 +1,1036 @@ +"""Plotly rendering backend: interactive 2-D surfaces and 3-D volumes. + +Imports only ``gdatastate``/``numerics`` (plus Plotly/Matplotlib themselves), +mirroring ``matplotlib.py``. Plotly cannot render mathtext, so labels go +through ``render.labels.latex_to_html`` instead. + +Like ``matplotlib.py``, ``plotly`` and ``plotly_animate`` own their *entire* +save/preview lifecycle here -- +``save``/``saveas``/``show`` (plus the rotating-export camera parameters) are +real parameters of both functions, so e.g. +``pg.load(f).interpolate().plotly(show=True)`` opens an auto-rotating browser +preview with zero CLI glue. Both default to inert +(``show=False``, ``save=False``): a bare call just builds and returns the +figure, so a script or unit test never has an unrequested file or browser +side effect. The generated ``plotly`` and ``plotly_animate`` commands lower +their canonical render signatures directly; there is no backend-specific +plotting command or operations wrapper. +""" + +from __future__ import annotations + +import os.path +import tempfile +import webbrowser +from pathlib import Path +from typing import Annotated + +import matplotlib as mpl +import numpy as np +import plotly.graph_objects as go +from plotly.subplots import make_subplots + +from postgkyl.cli_spec import ( + CliType, + CommandSpec, + Execution, + KeyValue, + PipelineInput, + ResultPolicy, + Section, + command, +) +from postgkyl.gdatastate import GDataState, materialize_point_values +from postgkyl.numerics import downsample, nodal_to_cell_centered_grid + +from ._ffmpeg import require_ffmpeg +from ._prep import resolve_axis_labels, squeeze_collapsed_axes, subplot_grid +from .labels import latex_to_html +from .style import DEFAULT_STYLE, apply_style + + +def _apply_plot_style(style: str | None, + rcParams: dict | None, + diverging: bool, + cmap: str | None, + xkcd: bool, + *, + background: str = "dark", + invert_cmap: bool = False) -> dict: + """Apply Matplotlib styling (colormap source) and return Plotly theme colors.""" + import matplotlib.pyplot as plt + + background_name = (background or "dark").strip().lower() + if style: + apply_style(style) + elif background_name == "light": + apply_style("default") + else: + apply_style(DEFAULT_STYLE) + + if background_name == "light": + mpl.rcParams["figure.facecolor"] = "#ffffff" + mpl.rcParams["axes.facecolor"] = "#ffffff" + mpl.rcParams["savefig.facecolor"] = "#ffffff" + mpl.rcParams["text.color"] = "#111111" + mpl.rcParams["axes.labelcolor"] = "#111111" + mpl.rcParams["xtick.color"] = "#111111" + mpl.rcParams["ytick.color"] = "#111111" + mpl.rcParams["axes.edgecolor"] = "#222222" + mpl.rcParams["grid.color"] = "#b8b8b8" + theme_colors = dict(paper_color="#ffffff", + scene_color="#ffffff", + text_color="#111111", + grid_color="#b8b8b8", + axis_line_color="#222222") + else: + theme_colors = dict(paper_color="#000000", + scene_color="#000000", + text_color="#e6e6e6", + grid_color="#2a3242", + axis_line_color="#9aa3b2") + + if rcParams: + for key, value in rcParams.items(): + mpl.rcParams[key] = value + + cmap_name = cmap if cmap is not None else ( + "RdBu_r" if diverging else "inferno") + mpl.rcParams["image.cmap"] = cmap_name + + if invert_cmap: + current = mpl.rcParams["image.cmap"] + mpl.rcParams["image.cmap"] = (current[:-2] + if current.endswith("_r") else f"{current}_r") + + if xkcd: + plt.xkcd() + + return theme_colors + + +def _plotly_colorscale(cmap_name: str, n: int = 256): + """Convert a Matplotlib colormap to a Plotly colorscale.""" + cmap = mpl.colormaps.get_cmap(cmap_name).resampled(n) + xs = np.linspace(0.0, 1.0, n) + colorscale = [] + for x, rgba in zip(xs, cmap(xs)): + r, g, b, a = rgba + colorscale.append([ + float(x), + f"rgba({int(r * 255)}, {int(g * 255)}, {int(b * 255)}, {float(a):.3f})" + ]) + return colorscale + + +def _opacity_mapping(colorscale, + min_alpha: float, + max_alpha: float, + log_scale: bool = False): + """Remap a Plotly colorscale's alpha channel over ``[min_alpha, max_alpha]``.""" + min_a = float(np.clip(min_alpha, 0.0, 1.0)) + max_a = float(np.clip(max_alpha, 0.0, 1.0)) + if max_a < min_a: + min_a, max_a = max_a, min_a + + out = [] + for stop, color in colorscale: + stop_value = float(stop) + mapped_stop = (np.log10(1.0 + 99.0 * stop_value) / + np.log10(100.0) if log_scale else stop_value) + if isinstance(color, + str) and color.startswith("rgba(") and color.endswith(")"): + parts = [part.strip() for part in color[5:-1].split(",")] + if len(parts) == 4: + r, g, b = parts[0], parts[1], parts[2] + alpha = min_a + (max_a - min_a) * mapped_stop + out.append([stop_value, f"rgba({r}, {g}, {b}, {alpha:.3f})"]) + else: + out.append([stop_value, color]) + else: + out.append([stop_value, color]) + return out + + +def _finite_range(values: np.ndarray) -> tuple[float, float]: + """Finite min/max of an array, ignoring NaN/inf.""" + finite = np.isfinite(values) + if np.any(finite): + finite_values = values[finite] + return float(np.nanmin(finite_values)), float(np.nanmax(finite_values)) + return float("nan"), float("nan") + + +def _axis_range(values: np.ndarray, axis_range, log_axis: bool = False): + """Axis range for a colorbar or scene axis, log10'd when ``log_axis``.""" + lower, upper = _finite_range(values) if axis_range is None else axis_range + if log_axis: + lower, upper = np.log10(lower), np.log10(upper) + return [lower, upper] + + +def _log_colorbar_ticks(log_min: float, log_max: float, max_ticks: int = 7): + """Tick values/text for a logarithmic (decade) colorbar.""" + if not np.isfinite(log_min) or not np.isfinite(log_max): + return [], [] + lo, hi = int(np.floor(log_min)), int(np.ceil(log_max)) + hi = max(hi, lo) + count = hi - lo + 1 + step = max(1, int(np.ceil(count / max_ticks))) + tick_vals = list(range(lo, hi + 1, step)) + if tick_vals[-1] != hi: + tick_vals.append(hi) + return [float(v) + for v in tick_vals], [f"10{v:d}" for v in tick_vals] + + +def _apply_log_colorscale(render_color_value: np.ndarray, cmin_val, cmax_val, + colorbar_kwargs: dict): + """Map color values into log10 space; adds decade tick config in place.""" + log_value = np.full(render_color_value.shape, np.nan, dtype=float) + valid_mask = render_color_value > 0 + log_value[valid_mask] = np.log10(render_color_value[valid_mask]) + + if np.any(valid_mask): + valid_min = float(np.nanmin(log_value[valid_mask])) + valid_max = float(np.nanmax(log_value[valid_mask])) + else: + valid_min, valid_max = 0.0, 1.0 + + if cmin_val is not None and cmin_val > 0: + valid_min = float(np.log10(cmin_val)) + if cmax_val is not None and cmax_val > 0: + valid_max = float(np.log10(cmax_val)) + if not np.isfinite(valid_max) or valid_max <= valid_min: + valid_max = valid_min + 1.0 + + render_color_value = np.nan_to_num(log_value, + nan=valid_min, + posinf=valid_max, + neginf=valid_min) + + tick_vals, tick_text = _log_colorbar_ticks(valid_min, valid_max) + if tick_vals: + colorbar_kwargs["tickmode"] = "array" + colorbar_kwargs["tickvals"] = tick_vals + colorbar_kwargs["ticktext"] = tick_text + return render_color_value, valid_min, valid_max + + +def _resolve_plotly_aspect(aspect: str | float | None): + """Resolve ``aspectmode``/``aspectratio`` for a Plotly 3-D scene.""" + if aspect is None: + return "auto", None + if isinstance(aspect, str): + aspect_value = aspect.strip().lower() + if aspect_value in ("auto", "data", "cube"): + return aspect_value, None + ratio = float(aspect) + return "manual", dict(x=ratio, y=ratio, z=ratio) + ratio = float(aspect) + return "manual", dict(x=ratio, y=ratio, z=ratio) + + +def _build_rotation_post_script(scene_name: str, + starting_azimuthal_angle: float, + polar_angle: float, rotation_period: float, + radius: float) -> str: + """Fill in the packaged rotation-controls JS template with camera params.""" + template_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + "rotation_controls.js") + with open(template_path) as template_file: + template = template_file.read() + replacements = { + "__PGKYL_SCENE_NAME__": scene_name, + "__PGKYL_AZIMUTH_DEG__": f"{float(starting_azimuthal_angle):.17g}", + "__PGKYL_POLAR_DEG__": f"{float(polar_angle):.17g}", + "__PGKYL_PERIOD_SEC__": f"{float(rotation_period):.17g}", + "__PGKYL_RADIUS__": f"{float(radius):.17g}", + } + for token, value in replacements.items(): + template = template.replace(token, value) + return template + + +def save_rotating_plotly_figure(fig, + file_name: str, + starting_azimuthal_angle: float, + fps: int, + polar_angle: float, + rotation_period: float, + radius: float = 2.0) -> None: + """Save a rotating Plotly 3-D figure as a GIF, MP4, or self-rotating HTML. + + Rotates the camera 360 degrees around the vertical axis, starting from + ``starting_azimuthal_angle`` degrees. ``.gif``/``.mp4`` render frame-by-frame + through Kaleido and ffmpeg; ``.html`` embeds a small JS animation loop + instead (no external process). + """ + import subprocess + + root, ext = os.path.splitext(file_name) + ext = ext.lower() + if ext not in (".gif", ".mp4", ".html"): + raise ValueError( + "save_rotating_plotly_figure expects an output ending with .gif, " + ".mp4, or .html") + if fps <= 0: + raise ValueError("fps must be a positive integer") + if rotation_period <= 0: + raise ValueError("rotation_period must be positive") + + scene_names = [ + name for name in fig.layout.to_plotly_json().keys() + if name == "scene" or name.startswith("scene") + ] + if not scene_names: + raise ValueError("Rotating export requires a Plotly 3D scene figure") + scene_name = scene_names[0] + + polar_rad = np.deg2rad(polar_angle) + xy_radius = radius * np.sin(polar_rad) + z_eye = radius * np.cos(polar_rad) + + if ext == ".html": + theta0 = np.deg2rad(starting_azimuthal_angle) + initial_camera = dict(eye=dict(x=float(xy_radius * np.cos(theta0)), + y=float(xy_radius * np.sin(theta0)), + z=float(z_eye)), + up=dict(x=0.0, y=0.0, z=1.0), + center=dict(x=0.0, y=0.0, z=0.0)) + fig.update_layout(**{scene_name: dict(camera=initial_camera)}) + + omega = 2.0 * np.pi / float(rotation_period) + if omega > 0.0: + post_script = _build_rotation_post_script(scene_name, + starting_azimuthal_angle, + polar_angle, rotation_period, + radius) + fig.write_html(file_name, include_plotlyjs="cdn", post_script=post_script) + else: + fig.write_html(file_name) + return + + with tempfile.TemporaryDirectory(prefix="pgkyl_rotate_") as tmp_dir: + frame_pattern = os.path.join(tmp_dir, "frame_%05d.png") + num_frames = max(2, int(round(float(fps) * float(rotation_period)))) + # Kaleido launches a fresh headless-Chrome process per to_image() call + # unless a persistent render server is running; for a multi-frame export + # that means one Chrome startup per frame. Hold the server open for the + # whole loop so only the first frame pays that cost. + import kaleido + kaleido.start_sync_server(silence_warnings=True) + try: + for idx in range(num_frames): + theta = np.deg2rad(starting_azimuthal_angle + 360.0 * idx / num_frames) + camera = dict(eye=dict(x=float(xy_radius * np.cos(theta)), + y=float(xy_radius * np.sin(theta)), + z=float(z_eye)), + up=dict(x=0.0, y=0.0, z=1.0), + center=dict(x=0.0, y=0.0, z=0.0)) + fig.update_layout(**{name: dict(camera=camera) for name in scene_names}) + png_bytes = fig.to_image(format="png") + with open(os.path.join(tmp_dir, f"frame_{idx:05d}.png"), + "wb") as frame_file: + frame_file.write(png_bytes) + finally: + kaleido.stop_sync_server(silence_warnings=True) + + ffmpeg_exe = require_ffmpeg("plotly_animate") + if ext == ".mp4": + ffmpeg_cmd = [ + ffmpeg_exe, "-y", "-framerate", + str(fps), "-i", frame_pattern, "-pix_fmt", "yuv420p", file_name + ] + else: + ffmpeg_cmd = [ + ffmpeg_exe, "-y", "-framerate", + str(fps), "-i", frame_pattern, "-vf", + "split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse", file_name + ] + subprocess.run(ffmpeg_cmd, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE) + + +def _default_output_stem(data: "GDataState") -> str: + """Best-effort output-file stem for a dataset with no explicit ``saveas``.""" + file_name = getattr(data, "_file_name", "") or "" + if file_name: + return os.path.basename(file_name).split(".")[0] + label = data.get_label() if hasattr(data, "get_label") else "" + return label or "plotly_output" + + +def _write_plotly_output(fig, file_name: str, *, + starting_azimuthal_angle: float, polar_angle: float, + rotation_period: float, fps: int) -> str: + """Save ``fig`` to ``file_name``, returning the (possibly extension-coerced) path. + + ``.mp4``/``.gif``/``.html`` rotate the camera on save (via + :func:`save_rotating_plotly_figure`); any other extension -- or none -- + is coerced to a plain, non-rotating ``.html``. + """ + root, ext = os.path.splitext(file_name) + ext = ext.lower() + if ext in (".mp4", ".gif", ".html"): + save_rotating_plotly_figure( + fig, + file_name, + starting_azimuthal_angle=starting_azimuthal_angle, + polar_angle=polar_angle, + rotation_period=rotation_period, + fps=fps) + return file_name + file_name = f"{root}.html" if root else f"{file_name}.html" + fig.write_html(file_name) + return file_name + + +def _preview_plotly_figure(fig, base_name: str, *, + starting_azimuthal_angle: float, polar_angle: float, + rotation_period: float, fps: int) -> str: + """Write a temp, auto-rotating HTML preview of ``fig`` and return its path.""" + safe_base = "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" + for ch in base_name).strip("_") + if not safe_base: + safe_base = "plotly_preview" + file_name = os.path.join(tempfile.gettempdir(), f"{safe_base}_preview.html") + save_rotating_plotly_figure(fig, + file_name, + starting_azimuthal_angle=starting_azimuthal_angle, + polar_angle=polar_angle, + rotation_period=rotation_period, + fps=fps) + return file_name + + +def open_preview(path: str) -> None: + """Open a saved HTML file in the default web browser.""" + webbrowser.open(Path(path).resolve().as_uri()) + + +def _prepare_3d_coordinates(coords, value_shape): + arrays = tuple(np.asarray(coord) for coord in coords) + if len(arrays) != 3: + raise ValueError( + "Plotly 3D plotting requires exactly three coordinate arrays") + if all(array.ndim == 1 for array in arrays): + mesh = np.meshgrid(*arrays, indexing="ij") + return mesh[0], mesh[1], mesh[2] + return arrays[0], arrays[1], arrays[2] + + +def _prepare_2d_coordinates(coords, value_shape): + arrays = tuple(np.asarray(coord) for coord in coords) + if len(arrays) != 2: + raise ValueError( + "Plotly surface plotting requires exactly two coordinate arrays") + if all(array.ndim == 1 for array in arrays): + mesh = np.meshgrid(*arrays, indexing="ij") + return mesh[0], mesh[1] + return arrays[0], arrays[1] + + +def _scene_axis(label: str | None, log_axis: bool, axis_range, showgrid: bool, + theme: dict) -> dict: + """A themed Plotly 3-D scene axis dict, shared by the x/y/z axes.""" + return dict(title=dict(text=latex_to_html(label), + font=dict(color=theme["text_color"])), + showgrid=showgrid, + type="log" if log_axis else "linear", + exponentformat="e", + range=axis_range, + showbackground=True, + backgroundcolor=theme["scene_color"], + gridcolor=theme["grid_color"], + linecolor=theme["axis_line_color"], + tickfont=dict(color=theme["text_color"]), + zerolinecolor=theme["grid_color"]) + + +def plotly(data: GDataState, + *, + squeeze: bool = False, + num_subplot_row: int | None = None, + num_subplot_col: int | None = None, + scatter: bool = False, + marker_radius: float = 4.0, + markerstyle: str = "circle", + diverging: bool = False, + xscale: float = 1.0, + xshift: float = 0.0, + yscale: float = 1.0, + yshift: float = 0.0, + zscale: float = 1.0, + zshift: float = 0.0, + cmin: float | None = None, + cmax: float | None = None, + cscale: float = 1.0, + cshift: float = 0.0, + clim: tuple[float, float] | None = None, + style: str | None = None, + rcParams: Annotated[dict[str, object] | None, + CliType(dict[str, str] | None), + KeyValue()] = None, + background: str = "dark", + invert_cmap: bool = False, + no_legend: bool = False, + label_prefix: str = "", + no_colorbar: bool = False, + xlabel: str | None = None, + ylabel: str | None = None, + zlabel: str | None = None, + clabel: str | None = None, + title: str | None = None, + logx: bool = False, + logy: bool = False, + logz: bool = False, + logc: bool = False, + aspect: Annotated[str | float | None, + CliType(str | None)] = None, + no_showgrid: bool = False, + hashtag: bool = False, + xkcd: bool = False, + color: str | None = None, + opacity: float | None = 1.0, + scatter_opacity_range: tuple[float, float] | None = None, + scatter_opacity_log: bool = False, + maximum_points_per_axis: int = 0, + surface_count: int = 32, + xrange: tuple[float, float] | None = None, + yrange: tuple[float, float] | None = None, + zrange: tuple[float, float] | None = None, + figsize: tuple[int, int] | None = None, + cylindrical_to_cartesian: bool = False, + cmap: str | None = None, + save: bool = False, + saveas: str | None = None, + show: bool = False, + azimuthal_angle: float = 0.0, + polar_angle: float = 85.0, + rotation_period: float = 40.0, + fps: int = 1): + """Render 2-D surface or 3-D volumetric data with Plotly. + + 2-D data (``num_dims == 2``, after squeezing any size-1 axis) is drawn as + a ``go.Surface`` (height map); 3-D data is drawn as a ``go.Volume`` or, + with ``scatter=True``, a ``go.Scatter3d`` point cloud. Multi-component + data lays out one scene per component unless ``squeeze`` is set. + + ``save``/``saveas``/``show`` make this call self-sufficient without any CLI + glue: ``show=True`` opens an auto-rotating HTML preview in the browser; + ``saveas`` (or ``save=True`` for an auto-derived name from ``data``'s + source file) writes it instead -- ``.mp4``/``.gif``/``.html`` extensions + get the rotating camera baked in (via :func:`save_rotating_plotly_figure`), + any other extension a plain static ``.html``. If both a save and + ``show=True`` are requested, the just-saved file is what opens (no + separate preview render). All three default to inert (``show=False``, + ``save=False``) -- a bare ``pg.load(f).interpolate().plotly()`` just + builds and returns the figure, no file written and no browser opened; + the generated CLI has exactly the same defaults. + + Args: + data: Point-value dataset to render. + squeeze: Draw only the first component and collapse singleton axes. + num_subplot_row: Forced row count for multi-component subplot layouts. + num_subplot_col: Forced column count for multi-component subplot layouts. + scatter: Draw 3-D data as a point cloud instead of a volume. + marker_radius: Scatter-marker radius. + markerstyle: Plotly scatter-marker symbol. + diverging: Center a diverging color range on zero. + xscale: Horizontal-coordinate scale factor. + xshift: Horizontal-coordinate shift applied before scaling. + yscale: Vertical-coordinate scale factor. + yshift: Vertical-coordinate shift applied before scaling. + zscale: Third-coordinate scale, or 2-D surface-height scale. + zshift: Third-coordinate shift, or 2-D surface-height shift. + cmin: Color-range lower bound. + cmax: Color-range upper bound. + cscale: Color-value scale factor. + cshift: Color-value shift. + clim: Explicit ``(minimum, maximum)`` color range. + style: Postgkyl/Matplotlib style used to derive colors. + rcParams: Matplotlib configuration overrides used while deriving styles. + background: ``"dark"`` or ``"light"`` scene theme. + invert_cmap: Reverse the selected colormap. + no_legend: Suppress labeled traces in the legend. + label_prefix: Prefix for component trace names. + no_colorbar: Suppress the color bar. + xlabel: Horizontal-axis label override. + ylabel: Vertical-axis label override. + zlabel: Third-axis label override. + clabel: Color-bar label override. + title: Figure-title override. + logx: Use a logarithmic horizontal axis. + logy: Use a logarithmic vertical axis. + logz: Use a logarithmic third axis. + logc: Use logarithmic color values. + aspect: Scene aspect mode (``auto``, ``cube``, ``data``), or numeric ratio. + no_showgrid: Suppress scene grid lines. + hashtag: Add a ``#pgkyl`` annotation. + xkcd: Derive colors from Matplotlib's XKCD style. + color: Replace the colormap with one fixed trace color. + opacity: Surface, volume, or marker opacity. + scatter_opacity_range: Minimum and maximum opacity encoded in scatter colors. + scatter_opacity_log: Map scatter opacity logarithmically. + maximum_points_per_axis: Downsample 3-D data to this many points per axis; + zero disables downsampling. + surface_count: Number of isosurfaces used by volume rendering. + xrange: Explicit horizontal-axis range. + yrange: Explicit vertical-axis range. + zrange: Explicit third-axis range. + figsize: Figure width and height in hundreds of pixels. + cylindrical_to_cartesian: Interpret 3-D coordinates as ``(R, Z, phi)``. + cmap: Matplotlib colormap name. + save: Save using a name derived from the input dataset. + saveas: Explicit output path. + show: Open the saved or temporary HTML preview in a browser. + azimuthal_angle: Initial camera azimuth for rotating output. + polar_angle: Camera polar angle for rotating output. + rotation_period: Seconds per camera revolution for animated output. + fps: Frames per second for animated output. + + Returns: + plotly.graph_objects.Figure: the assembled figure. + """ + data = materialize_point_values(data) + theme_colors = _apply_plot_style(style, + rcParams, + diverging, + cmap, + xkcd, + background=background, + invert_cmap=invert_cmap) + + grid, values = squeeze_collapsed_axes(list(data.grid), data.values) + num_dims = len(grid) + surface_mode = (num_dims == 2) + if num_dims not in (2, 3): + raise ValueError( + "plotly handles only 2D surface data or 3D volumetric data") + if surface_mode and scatter: + raise ValueError("Surface plots do not support scatter mode") + + # In surface mode the vertical axis is the function value, not a + # coordinate; default its label to empty unless the caller overrode it. + if surface_mode and zlabel is None: + zlabel = " " + xlabel, ylabel, zlabel, clabel = resolve_axis_labels(xlabel=xlabel, + ylabel=ylabel, + zlabel=zlabel, + clabel=clabel or "", + num_dims=num_dims, + xshift=xshift, + yshift=yshift, + zshift=zshift, + xscale=xscale, + yscale=yscale, + zscale=zscale) + + num_comps = values.shape[-1] + idx_comps = range(num_comps) + + if squeeze or num_comps == 1: + fig = go.Figure() + scene_names = ["scene"] + grid_shape = (1, 1) + else: + num_rows, num_cols = subplot_grid(num_comps, num_subplot_row, + num_subplot_col) + specs = [[{ + "type": "scene" + } for _ in range(num_cols)] for _ in range(num_rows)] + fig = make_subplots(rows=num_rows, cols=num_cols, specs=specs) + scene_names = [ + "scene" if idx == 0 else f"scene{idx + 1}" for idx in range(num_comps) + ] + grid_shape = (num_rows, num_cols) + + colorscale = _plotly_colorscale(mpl.rcParams["image.cmap"]) + scalar_colorscale = [[0.0, color], [1.0, color] + ] if bool(color) else colorscale + paper_color = theme_colors["paper_color"] + scene_color = theme_colors["scene_color"] + text_color = theme_colors["text_color"] + + fig.update_layout(paper_bgcolor=paper_color, + plot_bgcolor=paper_color, + font=dict(color=text_color)) + + colorbar_kwargs = dict(title=dict(text=clabel or "", + font=dict(color=text_color)), + exponentformat="e", + showexponent="all", + tickfont=dict(color=text_color), + bgcolor=paper_color) + + for comp_idx, comp in enumerate(idx_comps): + if comp_idx >= len(scene_names): + break + scene_name = scene_names[comp_idx] + row = 1 if grid_shape == (1, 1) else int(comp_idx / grid_shape[1]) + 1 + col = 1 if grid_shape == (1, 1) else int(comp_idx % grid_shape[1]) + 1 + label = f"{label_prefix:s}_c{comp:d}".strip("_") if len( + idx_comps) > 1 else label_prefix + cc_grid = nodal_to_cell_centered_grid(grid, values.shape[:num_dims]) + value = np.asarray(values[..., comp]) * zscale + zshift + color_value = value * cscale + cshift + render_color_value = np.array(color_value, copy=True) + value_min, value_max = _finite_range(color_value) + + if surface_mode: + x_grid, y_grid = _prepare_2d_coordinates(cc_grid, value.shape) + x = (np.asarray(x_grid) + xshift) * xscale + y = (np.asarray(y_grid) + yshift) * yscale + z = np.asarray(value) + else: + x_grid, y_grid, z_grid = _prepare_3d_coordinates(cc_grid, value.shape) + x_coord, y_coord, z_coord = (np.asarray(x_grid), np.asarray(y_grid), + np.asarray(z_grid)) + if cylindrical_to_cartesian: + # mapc2p cylindrical ordering is (R, Z, phi) + r, z_cyl, phi = x_coord, y_coord, z_grid + x_coord = r * np.cos(phi) + y_coord = r * np.sin(phi) + z_coord = z_cyl + x = (x_coord + xshift) * xscale + y = (y_coord + yshift) * yscale + z = (z_coord + zshift) * zscale + x_axis_range = _axis_range(x, xrange, logx) + y_axis_range = _axis_range(y, yrange, logy) + z_axis_range = _axis_range(z, zrange, logz) + + scene_aspectmode, scene_aspectratio = _resolve_plotly_aspect(aspect) + scene = dict(xaxis=_scene_axis(xlabel, logx, x_axis_range, not no_showgrid, + theme_colors), + yaxis=_scene_axis(ylabel, logy, y_axis_range, not no_showgrid, + theme_colors), + zaxis=_scene_axis(zlabel, logz, z_axis_range, not no_showgrid, + theme_colors), + bgcolor=scene_color, + aspectmode=scene_aspectmode, + aspectratio=scene_aspectratio) + fig.update_layout(**{scene_name: scene}) + + if diverging: + cmax_val = float(np.nanmax(np.abs(color_value))) + cmin_val = -cmax_val + else: + if clim is not None: + cmin_local, cmax_local = clim + else: + cmin_local, cmax_local = cmin, cmax + cmin_val = cmin_local if cmin_local is not None else value_min + cmax_val = cmax_local if cmax_local is not None else value_max + + trace_colorscale = scalar_colorscale + trace_colorbar_kwargs = dict(colorbar_kwargs) + show_colorbar = not no_colorbar and comp_idx == 0 and not bool(color) + trace_name = label or f"c{comp}" + show_trace_legend = not no_legend and bool(label) + + if surface_mode: + if logc: + render_color_value, cmin_val, cmax_val = _apply_log_colorscale( + render_color_value, cmin_val, cmax_val, trace_colorbar_kwargs) + trace_list = [ + go.Surface(x=x, + y=y, + z=z, + surfacecolor=render_color_value, + colorscale=trace_colorscale, + cmin=cmin_val, + cmax=cmax_val, + showscale=show_colorbar, + colorbar=trace_colorbar_kwargs if show_colorbar else None, + opacity=opacity, + name=trace_name, + showlegend=show_trace_legend) + ] + else: + if logz: + positive = np.where(render_color_value > 0, render_color_value, np.nan) + render_color_value = np.log10(positive) + if cmin_val is not None: + cmin_val = np.log10(max(cmin_val, np.finfo(float).tiny)) + if cmax_val is not None: + cmax_val = np.log10(cmax_val) + if logc: + render_color_value, cmin_val, cmax_val = _apply_log_colorscale( + render_color_value, cmin_val, cmax_val, trace_colorbar_kwargs) + render_x, render_y, render_z, render_color_value = downsample( + x, + y, + z, + render_color_value, + maximum_points_per_axis=maximum_points_per_axis) + + if scatter: + marker_size = max(1.0, 2.0 * float(marker_radius)) + scatter_colorscale = trace_colorscale + scatter_opacity = opacity + if not bool(color) and scatter_opacity_range is not None: + min_alpha, max_alpha = scatter_opacity_range + scatter_colorscale = _opacity_mapping(trace_colorscale, + min_alpha=min_alpha, + max_alpha=max_alpha, + log_scale=scatter_opacity_log) + scatter_opacity = 1.0 + trace_list = [ + go.Scatter3d( + x=render_x.ravel(), + y=render_y.ravel(), + z=render_z.ravel(), + mode="markers", + marker=dict( + size=marker_size, + symbol=markerstyle, + color=render_color_value.ravel(), + colorscale=scatter_colorscale, + cmin=cmin_val, + cmax=cmax_val, + opacity=scatter_opacity, + showscale=show_colorbar, + colorbar=trace_colorbar_kwargs if show_colorbar else None), + name=trace_name, + showlegend=show_trace_legend) + ] + else: + volume_opacity_scale = [[0.0, 0.0], [0.5, 0.2], [1.0, 0.8]] + trace_list = [ + go.Volume(x=render_x.ravel(), + y=render_y.ravel(), + z=render_z.ravel(), + value=render_color_value.ravel(), + colorscale=trace_colorscale, + cmin=cmin_val, + cmax=cmax_val, + opacity=opacity, + opacityscale=volume_opacity_scale, + surface_count=surface_count, + showscale=show_colorbar, + colorbar=trace_colorbar_kwargs if show_colorbar else None, + name=trace_name, + showlegend=show_trace_legend) + ] + + for trace in trace_list: + if grid_shape == (1, 1): + fig.add_trace(trace) + else: + fig.add_trace(trace, row=row, col=col) + + if bool(title): + fig.update_layout(title=title) + if bool(hashtag): + fig.add_annotation(text="#pgkyl", + x=0.99, + y=0.01, + xref="paper", + yref="paper", + showarrow=False, + xanchor="right", + yanchor="bottom") + if bool(figsize): + fig.update_layout(width=figsize[0] * 100, height=figsize[1] * 100) + fig.update_layout(margin=dict(l=10, r=10, t=40 if title else 10, b=10)) + + output_path = None + if save or saveas: + output_path = _write_plotly_output(fig, + saveas or _default_output_stem(data), + starting_azimuthal_angle=azimuthal_angle, + polar_angle=polar_angle, + rotation_period=rotation_period, + fps=fps) + if show: + if output_path is None: + output_path = _preview_plotly_figure( + fig, + _default_output_stem(data), + starting_azimuthal_angle=azimuthal_angle, + polar_angle=polar_angle, + rotation_period=rotation_period, + fps=fps) + open_preview(output_path) + return fig + + +command( + CommandSpec(Section.RENDER, + Execution.TERMINAL_EACH, + result=ResultPolicy.SILENT))(plotly) + + +@command( + CommandSpec(Section.RENDER, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) +def plotly_animate(data: Annotated[list[GDataState], + PipelineInput()], + *, + frame_labels: list[str] | None = None, + frame_duration: int = 50, + transition_duration: int = 0, + from_start: bool = False, + no_redraw: bool = False, + save: bool = False, + saveas: str | None = None, + show: bool = False): + """Build a Plotly animation figure from a sequence of datasets. + + Renders the first dataset with :func:`plotly` to create the base figure, + then renders every subsequent dataset as an animation frame, wiring up + Play/Pause buttons and a frame slider. All datasets must produce the same + number of traces. + + Like :func:`plotly`, ``save``/``saveas``/``show`` are self-sufficient here + with zero CLI glue: ``show=True`` opens the animation (a plain HTML preview + -- the frame slider is its own scrubber, so unlike :func:`plotly` this + never rotates the camera on save) in the browser; both default to inert, + so a bare call just builds and returns the figure (see :func:`plotly`'s + docstring for why). The per-frame :func:`plotly` calls below always render + with ``save=False, show=False``: only *this* function's own save/preview, + on the assembled animation, should ever hit disk or a browser tab. + + Args: + data: Selected datasets, one per animation frame. + frame_labels: Optional label for each frame and slider step. + frame_duration: Duration of each animation frame in milliseconds. + transition_duration: Duration of transitions between frames. + from_start: Start playback from the first rather than current slider frame. + no_redraw: Do not redraw traces between frames. + save: Save the animation as ``plotly_animate.html``. + saveas: Explicit HTML output path. + show: Open the animation in a browser. + + Returns: + plotly.graph_objects.Figure: the assembled animation. + """ + data_sequence = list(data) + if not data_sequence: + raise ValueError("plotly_animate requires at least one dataset") + + base_fig = plotly(data_sequence[0], save=False, saveas=None, show=False) + num_traces = len(base_fig.data) + + if frame_labels is None: + frame_labels = [str(idx) for idx in range(len(data_sequence))] + if len(frame_labels) != len(data_sequence): + raise ValueError("frame_labels length must match data_sequence length") + + frames = [] + for idx, dat in enumerate(data_sequence): + if idx == 0: + continue + frame_fig = plotly(dat, save=False, saveas=None, show=False) + if len(frame_fig.data) != num_traces: + raise ValueError( + "All animation frames must produce the same number of traces; " + f"frame 0 has {num_traces:d}, frame {idx:d} has {len(frame_fig.data):d}." + ) + frames.append( + go.Frame(name=str(frame_labels[idx]), + data=list(frame_fig.data), + traces=list(range(num_traces)))) + + base_fig.frames = frames + + animation_args = { + "frame": { + "duration": int(frame_duration), + "redraw": not no_redraw + }, + "transition": { + "duration": int(transition_duration) + }, + "fromcurrent": not from_start + } + pause_args = { + "frame": { + "duration": 0, + "redraw": not no_redraw + }, + "transition": { + "duration": 0 + }, + "mode": "immediate" + } + + slider_steps = [{ + "label": + str(label), + "method": + "animate", + "args": [[str(label)], { + "mode": "immediate", + "frame": { + "duration": int(frame_duration), + "redraw": not no_redraw + }, + "transition": { + "duration": int(transition_duration) + } + }], + } for label in frame_labels] + + base_fig.update_layout( + updatemenus=[{ + "type": + "buttons", + "showactive": + False, + "buttons": [ + { + "label": "Play", + "method": "animate", + "args": [None, animation_args] + }, + { + "label": "Pause", + "method": "animate", + "args": [[None], pause_args] + }, + ], + "x": + 0.02, + "y": + 0.0, + "xanchor": + "left", + "yanchor": + "bottom", + }], + sliders=[{ + "active": 0, + "currentvalue": { + "prefix": "Frame: " + }, + "pad": { + "t": 24 + }, + "steps": slider_steps + }], + ) + + output_path = None + if save or saveas: + out_name = saveas or "plotly_animate.html" + if not str(out_name).lower().endswith(".html"): + out_name = f"{out_name}.html" + base_fig.write_html(out_name) + output_path = out_name + if show: + if output_path is None: + output_path = os.path.join(tempfile.gettempdir(), + "plotly-animate_preview.html") + base_fig.write_html(output_path) + open_preview(output_path) + return base_fig + + +__all__ = [ + "plotly", "plotly_animate", "save_rotating_plotly_figure", "open_preview" +] diff --git a/src/postgkyl/output/postgkyl.mplstyle b/src/postgkyl/render/postgkyl.mplstyle similarity index 92% rename from src/postgkyl/output/postgkyl.mplstyle rename to src/postgkyl/render/postgkyl.mplstyle index 69b58c44..ad7d784e 100644 --- a/src/postgkyl/output/postgkyl.mplstyle +++ b/src/postgkyl/render/postgkyl.mplstyle @@ -9,4 +9,4 @@ image.cmap : inferno image.origin : lower grid.linewidth : 0.5 grid.linestyle : : -axes.prop_cycle : cycler('color', [(0, 0.4470, 0.7410), (0.8500, 0.3250, 0.0980), (0.9290, 0.6940, 0.1250), (0.4940, 0.1840, 0.5560), (0.4660, 0.6740, 0.1880), (0.3010, 0.7450, 0.9330), (0.6350, 0.0780, 0.1840)]) \ No newline at end of file +axes.prop_cycle : cycler('color', [(0, 0.4470, 0.7410), (0.8500, 0.3250, 0.0980), (0.9290, 0.6940, 0.1250), (0.4940, 0.1840, 0.5560), (0.4660, 0.6740, 0.1880), (0.3010, 0.7450, 0.9330), (0.6350, 0.0780, 0.1840)]) diff --git a/src/postgkyl/render/pyvista.py b/src/postgkyl/render/pyvista.py new file mode 100644 index 00000000..1aba432e --- /dev/null +++ b/src/postgkyl/render/pyvista.py @@ -0,0 +1,358 @@ +"""Canonical PyVista renderer for 3-D scalar-field volumes and isosurfaces. + +``pg.pyvista``, ``GData.pyvista``, ``operations.pyvista``, and the generated +CLI are aliases or lowerings of the one public function in this module. +PyVista needs a working (possibly software/off-screen) OpenGL context; every +entry point re-raises a ``RuntimeError`` naming that requirement instead of +letting a VTK error surface from deep inside the library. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pyvista as pv + +from postgkyl.cli_spec import ( + CommandSpec, + Execution, + ResultPolicy, + Section, + command, +) +from postgkyl.gdatastate import GDataState, materialize_point_values +from postgkyl.numerics import downsample, nodal_to_cell_centered_grid + +from ._prep import resolve_axis_labels, squeeze_collapsed_axes +from .labels import latex_to_unicode + + +def _require_gl_context(action): + """Run ``action`` (a zero-arg callable), turning a VTK/GL failure into a + clear ``RuntimeError`` instead of an opaque one from deep inside VTK.""" + try: + return action() + except (RuntimeError, ValueError): + raise + except Exception as exc: # pragma: no cover - depends on the host's GL stack + raise RuntimeError( + "pyvista rendering requires a working (possibly off-screen) OpenGL " + f"context; the render backend raised: {exc!r}") from exc + + +@command( + CommandSpec(Section.RENDER, + Execution.TERMINAL_EACH, + result=ResultPolicy.SILENT)) +def pyvista(data: GDataState, + *, + no_show: bool = False, + no_spin: bool = False, + max_points_per_axis: int = -1, + contour_levels: int = 10, + is_log: bool = False, + volume: bool = False, + is_shaded: bool = False, + hide_axes: bool = False, + mesh_clip_plane: bool = False, + mesh_slice_plane: bool = False, + volume_clip_plane: bool = False, + cmin: float | None = None, + cmax: float | None = None, + aspect_ratio: tuple[float, float, float] = (1, 1, 1), + camera_azimuth: float = 0.0, + camera_elevation: float = -30.0, + opacity: str = "sigmoid_4", + cmap: str = "inferno", + xlabel: str | None = None, + ylabel: str | None = None, + zlabel: str | None = None, + clabel: str = "", + title: str | None = "", + diverging: bool = False, + cylindrical_to_cartesian: bool = False, + theme: str = "default", + saveas: str = "", + xscale: float = 1.0, + yscale: float = 1.0, + zscale: float = 1.0, + xshift: float = 0.0, + yshift: float = 0.0, + zshift: float = 0.0, + hide_zeros: bool = False): + """Render a 3-D scalar field with PyVista. + + Builds a structured grid from the (single-component) scalar values and + renders it as a volume, contour isosurfaces, or an interactive clip/slice + plane. The grid is normalized to ``aspect_ratio`` because PyVista handles + non-integer axis extents poorly. Only the first value component is used. + + Args: + data: dataset to plot; must be 3-D (after squeezing any size-1 axis). + no_show: Do not open an interactive render window; render off-screen. + no_spin: Do not auto-rotate the camera in interactive windows. + max_points_per_axis: downsample to at most this many points per axis; + ``-1`` disables downsampling. + contour_levels: Number of isosurfaces extracted unless ``volume`` is set. + is_log: color by log10 of the scalar (non-positive values masked). + volume: Render a volume instead of isosurface contours. + is_shaded: enable shading on the volume render (volume mode only). + hide_axes: hide the bounding-box axes and labels. + mesh_clip_plane: add an interactive clip plane along ``-x``. + mesh_slice_plane: add an interactive slice plane along ``-x``. + volume_clip_plane: add an interactive volume clip plane (volume mode). + cmin: Color-limit lower bound; defaults to the data minimum. + cmax: Color-limit upper bound; defaults to the data maximum. + aspect_ratio: per-axis aspect the grid is normalized to. + camera_azimuth: Initial camera azimuth in degrees. + camera_elevation: Initial camera elevation in degrees. + opacity: a PyVista opacity preset string, ``"diverging"`` (opaque at + both ends, transparent in the middle), or a scalar opacity. + cmap: colormap name; overridden to ``"RdBu_r"`` when ``diverging``. + xlabel: Horizontal-axis label; auto-derived when omitted. + ylabel: Vertical-axis label; auto-derived when omitted. + zlabel: Third-axis label; auto-derived when omitted. + clabel: colorbar (scalar bar) title. + title: text drawn at the top of the render; omitted when ``None``. + diverging: use the diverging ``"RdBu_r"`` colormap. + cylindrical_to_cartesian: treat grid coordinates as cylindrical + ``(R, Z, phi)`` and convert to Cartesian before building the mesh. + theme: PyVista plot theme name; ``"default"`` leaves it unchanged. + saveas: output path; extension selects the exporter (``.html``, + ``.png``/``.jpg``/``.jpeg``, ``.pdf``/``.svg``, ``.gltf``, ``.vtksz``). + Empty string disables saving. + xscale: Multiplicative scale recorded for the horizontal axis. + yscale: Multiplicative scale recorded for the vertical axis. + zscale: Multiplicative scale recorded for the third axis. + xshift: Additive shift applied to the horizontal axis. + yshift: Additive shift applied to the vertical axis. + zshift: Additive shift applied to the third axis. + hide_zeros: hide grid points whose scalar value is exactly zero. + + Returns: + None: the function renders and/or saves the plot for its side effects. + + Raises: + ValueError: ``data`` is not 3-D, or ``saveas`` has an unsupported + extension. + RuntimeError: PyVista could not obtain a working OpenGL context. + """ + data = materialize_point_values(data) + _valid_exts = ("", ".html", ".png", ".jpg", ".jpeg", ".pdf", ".svg", ".gltf", + ".vtksz") + if saveas and not os.path.splitext(saveas)[1]: + saveas += ".png" + if saveas != "" and not saveas.endswith(_valid_exts[1:]): + raise ValueError( + "Unsupported file format for saving. Supported formats are: " + ".html, .png, .jpg, .jpeg, .pdf, .svg, .gltf, .vtksz") + + grid, values = squeeze_collapsed_axes(list(data.grid), data.values) + num_dims = len(grid) + if num_dims != 3: + raise ValueError(f"pyvista renders 3D scalar fields only, got {num_dims}D") + xlabel, ylabel, zlabel, clabel = resolve_axis_labels(xlabel=xlabel, + ylabel=ylabel, + zlabel=zlabel, + clabel=clabel, + num_dims=num_dims, + xshift=xshift, + yshift=yshift, + zshift=zshift, + xscale=xscale, + yscale=yscale, + zscale=zscale) + + scalar = np.asarray(values[..., 0]) + x, y, z = nodal_to_cell_centered_grid(grid, scalar.shape, meshgrid=True) + if cylindrical_to_cartesian: + r, z_cyl, theta_ang = x, y, z + x = r * np.cos(theta_ang) + y = r * np.sin(theta_ang) + z = z_cyl + + xmax, xmin = np.max(x), np.min(x) + ymax, ymin = np.max(y), np.min(y) + zmax, zmin = np.max(z), np.min(z) + datamax, datamin = np.max(scalar), np.min(scalar) + x_range, y_range, z_range = xmax - xmin, ymax - ymin, zmax - zmin + + # Normalize to [-aspect, aspect] per axis -- PyVista struggles with + # non-integer axis extents. + x = (x - xmin) / x_range * aspect_ratio[0] * 2 - aspect_ratio[0] + y = (y - ymin) / y_range * aspect_ratio[1] * 2 - aspect_ratio[1] + z = (z - zmin) / z_range * aspect_ratio[2] * 2 - aspect_ratio[2] + + x, y, z, scalar = downsample(x, + y, + z, + scalar, + maximum_points_per_axis=max_points_per_axis) + + if diverging: + cmap = "RdBu_r" + if opacity == "diverging": + cx = np.linspace(0, 1, num=255) + opacity = np.abs(cx - 0.5) * 2 + + off_screen = saveas.endswith((".png", ".jpg", ".jpeg")) or no_show + + def _build_and_render(): + pl = pv.Plotter(window_size=(1400, 900), off_screen=off_screen) + grid3d = pv.StructuredGrid(x, y, z) + + if theme != "default": + pv.set_plot_theme(theme) + + if hide_zeros: + x_ind, y_ind, z_ind = np.where(scalar == 0) + zero_indices = np.ravel_multi_index((x_ind, y_ind, z_ind), + dims=scalar.shape, + order="F") + if zero_indices.size: + grid3d.hide_points(zero_indices) + + grid3d["f_raw"] = scalar.ravel(order="F") + field = np.asarray(grid3d["f_raw"], dtype=float) + + colorbarformat = "%.2e" + clim = (cmin if cmin is not None else datamin, + cmax if cmax is not None else datamax) + if is_log: + positive_mask = np.asarray(grid3d["f_raw"]) > 0.0 + field = np.full(field.shape, np.nan, dtype=float) + field[positive_mask] = np.log10( + np.asarray(grid3d["f_raw"])[positive_mask]) + finite_field = field[np.isfinite(field)] + colorbarformat = "10^%.1f" + clim = ( + np.log10(cmin) if cmin is not None else float(np.min(finite_field)), + np.log10(cmax) if cmax is not None else float(np.max(finite_field))) + grid3d["f_plot"] = field + + scalar_bar_args = {"title": latex_to_unicode(clabel), "fmt": colorbarformat} + + if not volume: + contours = grid3d.contour(isosurfaces=contour_levels, scalars="f_plot") + if mesh_clip_plane: + pl.add_mesh_clip_plane(contours, + cmap=cmap, + clim=clim, + normal="-x", + opacity=opacity, + scalar_bar_args=scalar_bar_args, + factor=1.0) + elif mesh_slice_plane: + pl.add_mesh_slice(contours, + cmap=cmap, + clim=clim, + normal="-x", + opacity=opacity, + scalar_bar_args=scalar_bar_args, + factor=1.0) + else: + pl.add_mesh(contours, + cmap=cmap, + clim=clim, + opacity=opacity, + scalar_bar_args=scalar_bar_args) + else: + if mesh_clip_plane: + pl.add_mesh_clip_plane(grid3d, + scalars="f_plot", + cmap=cmap, + clim=clim, + opacity=opacity, + normal="-x", + scalar_bar_args=scalar_bar_args, + factor=1.0) + elif mesh_slice_plane: + pl.add_mesh_slice(grid3d, + scalars="f_plot", + cmap=cmap, + clim=clim, + opacity=opacity, + normal="-x", + scalar_bar_args=scalar_bar_args, + factor=1.0) + else: + vol = pl.add_volume(grid3d, + scalars="f_plot", + cmap=cmap, + clim=clim, + opacity=opacity, + shade=is_shaded, + scalar_bar_args=scalar_bar_args) + if volume_clip_plane: + pl.add_volume_clip_plane(vol, normal="-x") + + if title is not None: + pl.add_text(latex_to_unicode(f"{title}"), + position="upper_edge", + font_size=12) + + if hide_axes: + pl.hide_axes() + else: + # The mesh itself is normalized to +/-aspect_ratio (see above), so its + # own bounds carry no physical meaning; axes_ranges relabels the ticks + # with the true (shift/scale-adjusted) physical extent instead. + pv_bounds = pl.bounds + axes_ranges = (-(xmin + xshift) * xscale * pv_bounds.x_min, + (xmax + xshift) * xscale * pv_bounds.x_max, + -(ymin + yshift) * yscale * pv_bounds.y_min, + (ymax + yshift) * yscale * pv_bounds.y_max, + -(zmin + zshift) * zscale * pv_bounds.z_min, + (zmax + zshift) * zscale * pv_bounds.z_max) + pl.show_bounds(xtitle=latex_to_unicode(xlabel), + ytitle=latex_to_unicode(ylabel), + ztitle=latex_to_unicode(zlabel), + axes_ranges=axes_ranges, + n_xlabels=3, + n_ylabels=3, + n_zlabels=3, + grid="back", + location="origin", + all_edges=True, + use_3d_text=False, + fmt="%.2e") + + pl.camera.azimuth = camera_azimuth + pl.camera.elevation = camera_elevation + if not no_spin: + state = {"angle": camera_azimuth, "interacting": False} + + def _rotate(_step): + if state["interacting"]: + return + state["angle"] += 0.5 + pl.camera.azimuth = state["angle"] % 360 + + def _on_click(*_args): + state["interacting"] = True + + pl.add_timer_event(max_steps=99999999, duration=50, callback=_rotate) + pl.iren.add_observer("LeftButtonPressEvent", _on_click) + + if saveas != "": + if saveas.endswith(".html"): + pl.export_html(saveas) + elif saveas.endswith((".pdf", ".svg")): + pl.save_graphic(saveas) + elif saveas.endswith((".png", ".jpg", ".jpeg")): + pl.screenshot(saveas) + elif saveas.endswith(".gltf"): + pl.export_gltf(saveas) + elif saveas.endswith(".vtksz"): + pl.export_vtksz(saveas) + + if not no_show: + pl.show() + else: + pl.close() + + _require_gl_context(_build_and_render) + + +__all__ = ["pyvista"] diff --git a/src/postgkyl/render/rotation_controls.js b/src/postgkyl/render/rotation_controls.js new file mode 100644 index 00000000..29dcb45b --- /dev/null +++ b/src/postgkyl/render/rotation_controls.js @@ -0,0 +1,263 @@ +const gd = document.getElementById('{plot_id}'); +const sceneName = '__PGKYL_SCENE_NAME__'; +const defaultAzimuthDeg = __PGKYL_AZIMUTH_DEG__; +const defaultPolarDeg = __PGKYL_POLAR_DEG__; +const defaultPeriodSec = __PGKYL_PERIOD_SEC__; +const defaultRadius = __PGKYL_RADIUS__; +let rafId = null; +let startMs = null; + +let azimuthDeg = defaultAzimuthDeg; +let polarDeg = defaultPolarDeg; +let periodSec = defaultPeriodSec; +let cameraRadius = defaultRadius; + +let theta0 = 0.0; +let omega = 0.0; +let xyRadius = 0.0; +let zEye = 0.0; + +const clampPositive = (value, fallback) => (Number.isFinite(value) && value > 0.0 ? value : fallback); + +const recomputeRotationParams = () => { + const polarRad = polarDeg * Math.PI / 180.0; + theta0 = azimuthDeg * Math.PI / 180.0; + xyRadius = cameraRadius * Math.sin(polarRad); + zEye = cameraRadius * Math.cos(polarRad); + omega = 2.0 * Math.PI / periodSec; +}; + +const updateCamera = (theta) => { + const camera = { + eye: {x: xyRadius * Math.cos(theta), y: xyRadius * Math.sin(theta), z: zEye}, + up: {x: 0.0, y: 0.0, z: 1.0}, + center: {x: 0.0, y: 0.0, z: 0.0} + }; + Plotly.relayout(gd, { [sceneName + '.camera']: camera }); +}; + +const startRotation = () => { + if (rafId === null) { + rafId = requestAnimationFrame(animate); + } +}; + +const stopRotation = () => { + if (rafId !== null) { + cancelAnimationFrame(rafId); + rafId = null; + } +}; + +const resetRotation = () => { + startMs = null; + updateCamera(theta0); + startRotation(); +}; + +const parent = gd.parentNode; +if (parent) { + if (getComputedStyle(parent).position === 'static') { + parent.style.position = 'relative'; + } + + const controls = document.createElement('div'); + controls.style.position = 'absolute'; + controls.style.top = '12px'; + controls.style.left = '12px'; + controls.style.zIndex = '20'; + controls.style.background = 'rgba(255, 255, 255, 0.92)'; + controls.style.border = '1px solid #b7bec8'; + controls.style.borderRadius = '8px'; + controls.style.padding = '8px 10px'; + controls.style.fontFamily = 'sans-serif'; + controls.style.fontSize = '12px'; + controls.style.color = '#1f2933'; + controls.style.boxShadow = '0 2px 8px rgba(0, 0, 0, 0.18)'; + controls.style.display = 'grid'; + controls.style.gridTemplateColumns = 'auto auto'; + controls.style.gap = '6px 8px'; + controls.style.alignItems = 'center'; + controls.style.opacity = '0'; + controls.style.pointerEvents = 'none'; + controls.style.transition = 'opacity 120ms ease'; + + const showControlsButton = document.createElement('button'); + showControlsButton.type = 'button'; + showControlsButton.textContent = 'Show rotation controls'; + showControlsButton.style.position = 'absolute'; + showControlsButton.style.top = '12px'; + showControlsButton.style.left = '12px'; + showControlsButton.style.zIndex = '21'; + showControlsButton.style.fontSize = '12px'; + showControlsButton.style.padding = '4px 8px'; + showControlsButton.style.cursor = 'pointer'; + showControlsButton.style.opacity = '0'; + showControlsButton.style.pointerEvents = 'none'; + showControlsButton.style.transition = 'opacity 120ms ease'; + + const makeNumberInput = (value, min, step) => { + const input = document.createElement('input'); + input.type = 'number'; + input.value = String(value); + input.min = String(min); + input.step = String(step); + input.style.width = '86px'; + input.style.fontSize = '12px'; + return input; + }; + + const addRow = (labelText, inputEl) => { + const label = document.createElement('label'); + label.textContent = labelText; + controls.appendChild(label); + controls.appendChild(inputEl); + }; + + const periodInput = makeNumberInput(defaultPeriodSec, 0.001, 0.1); + const azimuthInput = makeNumberInput(defaultAzimuthDeg, -3600, 1); + const polarInput = makeNumberInput(defaultPolarDeg, -3600, 1); + const radiusInput = makeNumberInput(defaultRadius, 0.001, 0.1); + + addRow('Period (s)', periodInput); + addRow('Azimuth (deg)', azimuthInput); + addRow('Polar (deg)', polarInput); + addRow('Radius', radiusInput); + + const buttonWrap = document.createElement('div'); + buttonWrap.style.gridColumn = '1 / span 2'; + buttonWrap.style.display = 'flex'; + buttonWrap.style.gap = '8px'; + + const applyButton = document.createElement('button'); + applyButton.type = 'button'; + applyButton.textContent = 'Apply'; + + const stopButton = document.createElement('button'); + stopButton.type = 'button'; + stopButton.textContent = 'Stop rotation'; + + const hideButton = document.createElement('button'); + hideButton.type = 'button'; + hideButton.textContent = 'Hide controls'; + + for (const btn of [applyButton, stopButton, hideButton]) { + btn.style.fontSize = '12px'; + btn.style.padding = '3px 8px'; + btn.style.cursor = 'pointer'; + } + + let controlsCollapsed = true; + let hoverActive = false; + let hideTimer = null; + + const setControlsVisible = (visible) => { + controls.style.opacity = visible ? '1' : '0'; + controls.style.pointerEvents = visible ? 'auto' : 'none'; + }; + + const setShowButtonVisible = (visible) => { + showControlsButton.style.opacity = visible ? '1' : '0'; + showControlsButton.style.pointerEvents = visible ? 'auto' : 'none'; + }; + + const refreshControlsVisibility = () => { + if (!hoverActive) { + setControlsVisible(false); + setShowButtonVisible(false); + return; + } + if (controlsCollapsed) { + setControlsVisible(false); + setShowButtonVisible(true); + } else { + setControlsVisible(true); + setShowButtonVisible(false); + } + }; + + const clearHideTimer = () => { + if (hideTimer !== null) { + clearTimeout(hideTimer); + hideTimer = null; + } + }; + + const scheduleHide = () => { + clearHideTimer(); + hideTimer = setTimeout(() => { + hoverActive = false; + refreshControlsVisibility(); + }, 100); + }; + + const applyInputs = () => { + periodSec = clampPositive(parseFloat(periodInput.value), defaultPeriodSec); + cameraRadius = clampPositive(parseFloat(radiusInput.value), defaultRadius); + azimuthDeg = Number.isFinite(parseFloat(azimuthInput.value)) ? parseFloat(azimuthInput.value) : defaultAzimuthDeg; + polarDeg = Number.isFinite(parseFloat(polarInput.value)) ? parseFloat(polarInput.value) : defaultPolarDeg; + + periodInput.value = String(periodSec); + radiusInput.value = String(cameraRadius); + azimuthInput.value = String(azimuthDeg); + polarInput.value = String(polarDeg); + + recomputeRotationParams(); + resetRotation(); + }; + + applyButton.addEventListener('click', () => { + applyInputs(); + }); + + stopButton.addEventListener('click', () => { + stopRotation(); + }); + + hideButton.addEventListener('click', () => { + controlsCollapsed = true; + refreshControlsVisibility(); + }); + + showControlsButton.addEventListener('click', () => { + controlsCollapsed = false; + hoverActive = true; + refreshControlsVisibility(); + }); + + parent.addEventListener('mouseenter', () => { + hoverActive = true; + clearHideTimer(); + refreshControlsVisibility(); + }); + + parent.addEventListener('mouseleave', () => { + scheduleHide(); + }); + + buttonWrap.appendChild(applyButton); + buttonWrap.appendChild(stopButton); + buttonWrap.appendChild(hideButton); + controls.appendChild(buttonWrap); + parent.appendChild(controls); + parent.appendChild(showControlsButton); + refreshControlsVisibility(); +} + +gd.addEventListener('mousedown', stopRotation); +gd.addEventListener('wheel', stopRotation); +gd.addEventListener('touchstart', stopRotation); + +const animate = (timestamp) => { + if (startMs === null) { + startMs = timestamp; + } + const elapsedSeconds = (timestamp - startMs) / 1000.0; + const theta = theta0 + omega * elapsedSeconds; + updateCamera(theta); + rafId = requestAnimationFrame(animate); +}; + +recomputeRotationParams(); +updateCamera(theta0); +startRotation(); diff --git a/src/postgkyl/render/style.py b/src/postgkyl/render/style.py new file mode 100644 index 00000000..bc803fde --- /dev/null +++ b/src/postgkyl/render/style.py @@ -0,0 +1,47 @@ +"""Matplotlib style application -- the ``apply_style`` verb-adjacent helper. + +The old ``utils/load_style.py`` hand-parsed an ``.mplstyle`` file line by +line (with a special case for ``cycler(...)`` values) into a Typer context's +``rcParams`` dict. Matplotlib's own style-file parser already supports that +exact ``cycler(...)`` syntax (see ``postgkyl.mplstyle``'s ``axes.prop_cycle`` +line), so re-implementing a parser here would be a second, hand-maintained +copy of a fact Matplotlib already owns (DOCTRINE V). This module is a thin, +context-free wrapper: ``apply_style`` resolves the packaged default/name and +forwards to ``matplotlib.pyplot.style.use``. +""" + +from __future__ import annotations + +import os.path + +_STYLE_DIR = os.path.dirname(os.path.realpath(__file__)) + +# Names this package ships a style sheet for, resolved before falling through +# to Matplotlib's own named styles / arbitrary file paths. +_PACKAGED_STYLES = { + "postgkyl": os.path.join(_STYLE_DIR, "postgkyl.mplstyle"), +} + +DEFAULT_STYLE = "postgkyl" + + +def apply_style(path_or_name: str | None = None) -> None: + """Apply a Matplotlib style, mutating ``matplotlib.rcParams`` in place. + + Args: + path_or_name: A packaged style name (currently only ``"postgkyl"``), a + name Matplotlib recognizes (e.g. ``"dark_background"``), or a path to + an ``.mplstyle`` file. ``None`` applies the packaged Postgkyl default. + + This is the module's one documented effect: it mutates global Matplotlib + rc state (there is no other way to apply a style; see + ``matplotlib.pyplot.style.use``). + """ + import matplotlib.pyplot as plt + + name = path_or_name or DEFAULT_STYLE + target = _PACKAGED_STYLES.get(name, name) + plt.style.use(target) + + +__all__ = ["apply_style", "DEFAULT_STYLE"] diff --git a/src/postgkyl/tools/__init__.py b/src/postgkyl/tools/__init__.py deleted file mode 100644 index 04feb53a..00000000 --- a/src/postgkyl/tools/__init__.py +++ /dev/null @@ -1,68 +0,0 @@ -from .calculus import integrate - -# import parameter computation functions -from .params import get_magB -from .params import get_vt -from .params import get_vA -from .params import get_omegaC -from .params import get_omegaP -from .params import get_d -from .params import get_lambdaD -from .params import get_rho -from .params import get_beta - -# import primitive variable functions -from .prim_vars import get_density -from .prim_vars import get_vx -from .prim_vars import get_vy -from .prim_vars import get_vz -from .prim_vars import get_vi -from .prim_vars import get_pxx -from .prim_vars import get_pxy -from .prim_vars import get_pxz -from .prim_vars import get_pyy -from .prim_vars import get_pyz -from .prim_vars import get_pzz -from .prim_vars import get_pij -from .prim_vars import get_p -from .prim_vars import get_ke -from .prim_vars import get_temp -from .prim_vars import get_sound -from .prim_vars import get_mach -from .prim_vars import get_mhd_Bx -from .prim_vars import get_mhd_By -from .prim_vars import get_mhd_Bz -from .prim_vars import get_mhd_Bi -from .prim_vars import get_mhd_mag_p -from .prim_vars import get_mhd_p -from .prim_vars import get_mhd_temp -from .prim_vars import get_mhd_sound -from .prim_vars import get_mhd_mach - -from .pressure_diagnostics import get_p_par -from .pressure_diagnostics import get_gkyl_10m_p_par -from .pressure_diagnostics import get_p_perp -from .pressure_diagnostics import get_gkyl_10m_p_perp -from .pressure_diagnostics import get_agyro -from .pressure_diagnostics import get_gkyl_10m_agyro - -from .accumulate_current import accumulate_current -from .calc_enstrophy import calc_enstrophy -from .calc_ke_dke import calc_ke_dke -from .energetics import energetics -from .fft import fft -from .growth import exp2 -from .growth import fit_growth -from .init_polar import init_polar -from .mag_sq import mag_sq -from .parrotate import parrotate -from .perprotate import perprotate -from .polar_isotropic import polar_isotropic -from .rel_change import rel_change - -# import filters.py functions -from .filters import fft_filtering -from .filters import butter_filtering - -from .laguerre_compose import laguerre_compose -from .transform_frame import transform_frame \ No newline at end of file diff --git a/src/postgkyl/tools/accumulate_current.py b/src/postgkyl/tools/accumulate_current.py deleted file mode 100644 index 5508ec94..00000000 --- a/src/postgkyl/tools/accumulate_current.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Postgkyl module for accumulating current.""" - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -from postgkyl.utils import input_parser - -if TYPE_CHECKING: - from postgkyl import GData -#end - - -def accumulate_current(data: GData | Tuple[list, np.ndarray], qbym: bool = False, - overwrite=False, stack=False) -> Tuple[list, np.ndarray]: - """Computes the current from an arbitrary number of input species. - - Args: - data: GData or grid and values - input field - NOTE: These are GData objects which include metadata such as charge and mass - qbym: bool = False - optional input for multiplying by charge/mass ratio instead of just charge - NOTE: Should be true for fluid data - - XXX overwrite and stack need refactoring; see laguerre_compose.py - """ - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - grid, values = input_parser(data) - out = np.zeros_like(values) - factor = 0.0 - if qbym and data.mass and data.charge is not None: - factor = data.charge/data.mass - else: - factor = -1.0 - # end - out = factor*values - if overwrite: - data.push(grid, out) - return grid, out diff --git a/src/postgkyl/tools/calc_enstrophy.py b/src/postgkyl/tools/calc_enstrophy.py deleted file mode 100755 index 2c53d19c..00000000 --- a/src/postgkyl/tools/calc_enstrophy.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Postgkyl module for calculating enstrophy.""" - -import numpy as np - -import postgkyl - - -def calc_enstrophy(info_file, init_frame, final_frame): - """Calculates the enstrophy in 2D in the general and incompressible forms. - - Calculates the enstrophy in 2D in the general form (integral of the magnitude squared - of the curl of the velocity over the surface) and incompressible form (integral of the - magnitude of the gradient of velocity squared over the surface). - - Only for 3D also compares the two results and determines if incompressibility is - conserved. - - Args: - XXX will be filled up after refactoring - """ - - # get the matrices: rho, px, py, pz - frame = postgkyl.GData(f"{info_file}{str(init_frame)}.bp") - data = frame.values - grid = frame.grid - dx = grid[0][1] - grid[0][0] - dy = grid[1][1] - grid[1][0] - dz = grid[2][1] - grid[2][0] - - r = 0 - enstrophy = np.zeros((1, (final_frame - init_frame + 1))) - incom_enstrophy = enstrophy - incom_mag = np.zeros( - (len(data[:, 0, 0, 0]), len(data[0, :, 0, 0]), len(data[0, 0, :, 0])) - ) - - for i in range(init_frame, final_frame + 1): - frame = postgkyl.GData(f"{info_file}{i:d}.bp") - data = frame.values - - rho = data[..., 0] - px = data[..., 1] - py = data[..., 2] - pz = data[..., 3] - # calculate ux, uy, uz - u = px / rho - v = py / rho - w = pz / rho - # calculate the gradient - u_gradient = np.gradient(u, dx, dy, dz, edge_order=2) - v_gradient = np.gradient(v, dx, dy, dz, edge_order=2) - w_gradient = np.gradient(w, dx, dy, dz, edge_order=2) - A = [u_gradient, v_gradient, w_gradient] - A = np.array(A) - - u_x = np.array(u_gradient[0]) - u_y = np.array(u_gradient[1]) - u_z = np.array(u_gradient[2]) - v_x = np.array(v_gradient[0]) - v_y = np.array(v_gradient[1]) - v_z = np.array(v_gradient[2]) - w_x = np.array(w_gradient[0]) - w_y = np.array(w_gradient[1]) - w_z = np.array(w_gradient[2]) - - # find enstrophy in terms of curl magnitude squared integrand - curl_mag = (w_y - v_z)**2 + (u_z - w_x)**2 + (v_x - u_y)**2 - enstrophy[0, r] = np.sum(curl_mag, axis=(0, 1, 2))*dx*dy*dz - - # find incompressible enstrophy magnitude squared integrand - for c in range(0, (len(u[:, 0, 0]) - 1)): - for j in range(0, (len(u[0, :, 0]) - 1)): - for k in range(0, (len(u[0, 0, :]) - 1)): - incom_mag[c, j, k] = ( - np.trace(np.transpose(A[:, :, c, j, k])*A[:, :, c, j, k])*rho[c, j, k] - ) - incom_enstrophy[0, r] = np.sum(incom_mag, axis=(0, 1, 2))*dx*dy*dz - r += 1 - #end - - return enstrophy, incom_enstrophy diff --git a/src/postgkyl/tools/calc_ke_dke.py b/src/postgkyl/tools/calc_ke_dke.py deleted file mode 100755 index 05285ab4..00000000 --- a/src/postgkyl/tools/calc_ke_dke.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Postgkyl module for calculating total kinetic energy.""" - -from typing import Tuple -import numpy as np - -import postgkyl - - -def calc_ke_dke(root_file_name: str, init_frame: int, final_frame: int, dim: int, - vol: float, init_time: float, final_time: float) -> Tuple[np.ndarray, np.ndarray]: - """CalculateS all the total kinetic energy and the rate of dissipation of KE - - Args: - root_file_name: str - the name of the file before the numbers start - init_frame: int - is the first frame - final_frame: int - is the final frame - dim: int - gives the dimension of the simulation (2 = 2D, 3 = 3D) - vol: float - the volume of the grid - - Returns: - kinetic energy and dissipation of KE - """ - - # calculate integrated kinetic energy - ke = np.zeros((1, (final_frame - init_frame + 1))) - dEk = ke - f = postgkyl.GData(f"{root_file_name}{str(init_frame)}.bp") - grid = f.get_grid() - dx = grid[0][1] - grid[0][0] - dy = grid[1][1] - grid[1][0] - dt = (final_time - init_time + 1) / (final_frame - init_frame + 1) - r = 0 - - if dim == 3: - dz = grid[2][1] - grid[2][0] - else: # dim == 2: - dz = 1 - - for c in range(init_frame, final_frame + 1): - frame = postgkyl.GData(f"root_file_name{c:d}.bp") - data = frame.get_values() - rho = data[..., 0] - px = data[..., 1] - py = data[..., 2] - pz = data[..., 3] - - u = px / rho - v = py / rho - w = pz / rho - - e = rho * (u**2 + v**2 + w**2) - ke[0, r] = np.sum(e, axis=(0, 1, 2))*dx*dy*dz*vol - r += 1 - - r = 0 - for i in range(init_frame, final_frame - 1): - dEk[0, r] = -(ke[0, i + 1] - ke[0, i]) / dt - r += 1 - # end - - return ke, dEk diff --git a/src/postgkyl/tools/calculus.py b/src/postgkyl/tools/calculus.py deleted file mode 100644 index aca3e692..00000000 --- a/src/postgkyl/tools/calculus.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Postgkyl module for calculating integrals and derivatives.""" - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -if TYPE_CHECKING: - from postgkyl import GData -# end - - -def integrate(data: GData, axis: int | tuple | str, - overwrite=False, stack=False) -> Tuple[list, np.ndarray]: - """Integrates Gkeyll data. - - Currently simply uses the NumPy dot function. True, DG integration should be - implemented at some point. - - Args: - data: GData - axis: int, tuple or str - Specify axis to integrate over - - XXX overwrite and stack need refactoring; see laguerre_compose.py - """ - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - grid = list(data.grid) - values = np.copy(data.values) - - # Convert Python input to an input Numpy understands - if axis is not None: - if isinstance(axis, int): - axis = tuple([axis]) - elif isinstance(axis, tuple): - pass - elif isinstance(axis, str): - if len(axis.split(",")) > 1: - axes = axis.split(",") - axis = tuple([int(a) for a in axes]) - elif len(axis.split(":")) == 2: - bounds = axis.split(":") - # axis = np.zeros(bounds[1]-bounds[0], np.int) - # axis += int(bounds[0]) - axis = tuple(range(bounds[0], bounds[1])) - else: - axis = tuple([int(axis)]) - # end - else: - raise TypeError( - "'axis' needs to be integer, tuple, string of comma separated integers, or a slice ('int:int')" - ) - # end - else: - num_dims = data.get_num_dims() - axis = tuple(range(num_dims)) - # end - - # Get dz elements - dz = [] - for d, coord in enumerate(grid): - dz.append(coord[1:] - coord[:-1]) - if len(coord) > 1 and len(coord) == values.shape[d]: - dz[-1] = np.append(dz[-1], dz[-1][-1]) - # end - # end - - # Integration assuming values are cell centered averages - # Should work for nonuniform meshes - for ax in sorted(axis, reverse=True): - if len(grid[ax]) > 1: - values = np.moveaxis(values, ax, -1) - values = np.dot(values, dz[ax]) - else: - values = values.mean(axis=ax) - # end - # end - - for ax in sorted(axis): - grid[ax] = np.array([grid[ax].mean()]) - values = np.expand_dims(values, ax) - # end - - if overwrite: - data.push(grid, values) - - return grid, values - # end - - -def grad(): - ... - - -def div(): - ... - - -def curl(): - ... diff --git a/src/postgkyl/tools/energetics.py b/src/postgkyl/tools/energetics.py deleted file mode 100644 index 0cb18cf2..00000000 --- a/src/postgkyl/tools/energetics.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Postgkyl module for separating energy components.""" - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -from postgkyl.tools import get_p, get_ke, mag_sq -if TYPE_CHECKING: - from postgkyl import GData -# end - -def energetics(data_elc: GData, data_ion: GData, data_field: GData) -> Tuple[list, np.ndarray]: - """Function to separate components of the energy. - - Works for both species and EM fields and separates into constituent parts (species - energy -> thermal + kinetic, field energy -> electric + magnetic) - - Args: - data_elc: GData - input GData object for electrons - data_ion: GData - input GData object for ions - data_field: GData - input GData object for EM fields - - XXX overwrite and stack need refactoring; see laguerre_compose.py - - Notes: - Assumes two-species plasma - """ - # Grid is the same for each of the input objects - grid = data_field.get_grid() - values_field = data_field.get_values() - # Output array is a seven component field - # 1) Electron thermal - # 2) Electron kinetic - # 3) Ion thermal - # 4) Ion kinetic - # 5) Electric - # 6) Magnetic - # 7) Total - out = np.zeros(values_field[..., :7].shape) - - grid, pre = get_p(data_elc) - grid, kee = get_ke(data_elc) - grid, pri = get_p(data_ion) - grid, kei = get_ke(data_ion) - # Can compute magnitude squared of electric and magnetic fields with magsq diagnostic - grid, esq = mag_sq(data_field, coords="0:3") - grid, bsq = mag_sq(data_field, coords="3:6") - - out[..., 0] = np.squeeze(pre) - out[..., 1] = np.squeeze(kee) - out[..., 2] = np.squeeze(pri) - out[..., 3] = np.squeeze(kei) - out[..., 4] = np.squeeze(esq/2.0) - out[..., 5] = np.squeeze(bsq/2.0) - out[..., 6] = np.squeeze(pre + kee + pri + kei + esq/2.0 + bsq/2.0) - return grid, out diff --git a/src/postgkyl/tools/fft.py b/src/postgkyl/tools/fft.py deleted file mode 100644 index e2bbca63..00000000 --- a/src/postgkyl/tools/fft.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Postgkyl module for wrapping FFT.""" - -from __future__ import annotations - -import numpy as np -import scipy.fft -from typing import Tuple, TYPE_CHECKING - -from postgkyl.tools.init_polar import init_polar -from postgkyl.tools.polar_isotropic import polar_isotropic -if TYPE_CHECKING: - from postgkyl import GData -# end - -def fft(data: GData, psd: bool = False, iso: bool = False, - overwrite: bool = False, stack: bool = False) -> Tuple[np.ndarray, np.ndarray]: - """Postgkyl wrapper of scipy FFT. - - Args: - data: GData - psd: bool - Flag to calculate the Power Spectral Density - iso: bool - Flag to return isotropic spectra - - XXX overwrite and stack need refactoring; see laguerre_compose.py - """ - if stack: - overwrite = stack - # end - grid = data.get_grid() - values = data.get_values() - - # Remove dummy dimensions - num_dims = len(grid) - idx = [] - for d in range(num_dims): - if len(grid[d]) <= 2: - idx.append(d) - # end - # end - if idx: - #grid = np.delete(grid, idx) - [grid.pop(i) for i in idx[::-1]] - values = np.squeeze(values, tuple(idx)) - num_dims = len(grid) - # end - num_comps = data.get_num_comps() - if num_dims == 1: - N = len(grid[0]) - dx = grid[0][1] - grid[0][0] - freq = [scipy.fft.fftfreq(N, dx)] - ft_values = np.zeros(values.shape, "complex") - for comp in np.arange(num_comps): - ft_values[..., comp] = scipy.fft.fft(values[..., comp]) - # end - - if psd: - freq[0] = freq[0][:N//2] - ft_values = np.abs(ft_values[:N//2, :])**2 - # end - - if overwrite: - data.push(freq, ft_values) - else: - return freq, ft_values - # end - else: - N = np.zeros(3, dtype=int) - dx = np.zeros(3) - freq = [] - for i in range(0, num_dims): - N[i] = len(grid[i]) - dx[i] = grid[i][1] - grid[i][0] - freq.append(scipy.fft.fftfreq(N[i], dx[i])) - # end - ft_values = np.zeros(values.shape, "complex") - for comp in np.arange(num_comps): - ft_values[..., comp] = scipy.fft.fftn(values[..., comp]) - # end - if psd: - for i in range(0, num_dims): - freq[i] = freq[i][:N[i]//2] - if num_dims == 2: - ft_values = np.abs(ft_values[:N[0]//2, :N[1]//2, :])**2 - # If only 2D, append third dummy index for ease of logic - freq.append(0) - elif num_dims == 3: - ft_values = np.abs(ft_values[:N[0]//2, :N[1]//2, :N[2]//2, :])**2 - else: - raise ValueError("Only 1D, 2D, and 3D data are currently supported.") - # end - if iso: - nkpolar = int(np.sqrt(np.sum(N[:] ** 2))) - nkx = N[0]//2 - nky = N[1]//2 - nkz = N[2]//2 - kx = freq[0] - ky = freq[1] - kz = freq[2] - akp, nbin, polar_index, _ = init_polar(nkx, nky, nkz, kx, ky, kz, nkpolar) - fft_iso = np.zeros((nkpolar, num_comps)) - for comp in np.arange(num_comps): - fft_iso[:, comp] = polar_isotropic(nkpolar, nkx, nky, nkz, polar_index, - nbin, ft_values[..., comp], kx, ky, kz) - # end - # Return isotropic spectra and 1D isotropic ks - if overwrite: - data.push([akp], fft_iso) - return [akp], fft_iso - # end - # end - # end - - if overwrite and not iso: - data.push(freq, ft_values) - return freq, ft_values - # end - # end diff --git a/src/postgkyl/tools/filters.py b/src/postgkyl/tools/filters.py deleted file mode 100644 index 3174fb4e..00000000 --- a/src/postgkyl/tools/filters.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Postgkyl module for filtering. - -Contains FFT and butter filters. -""" - -from scipy.signal import butter, lfilter -from typing import Optional -import matplotlib.pyplot as plt -import numpy as np - - -def _click_coords(event): - global ix, iy - ix, iy = event.xdata, event.ydata - plt.close() - - -def fft_filtering(data: np.ndarray, dt: float = 1.0, cutoff: Optional[float] = None) -> np.ndarray: - """Filter data using numpy FFT. - - Args: - data: np.ndarray - dt: float = 1.0 - set spacing of data - cutoff: float - set high frequency cut-off (default: None) - - Note: - If the cutoff is not selected, and interactive figure will pop out - that allows for the cut-off selection. - """ - N = len(data) - freq = np.fft.fftfreq(N, dt) - FT = np.fft.fft(data) - - # Get the cut-off frequency if not specified - if cutoff is None: - fig, ax = plt.subplots(1, 1) - # plot just N/2 points - ax.semilogy(freq[1:N//2], 2.0/N*np.abs(FT[1:N//2])) - ax.grid() - ax.set_xlabel("Freq") - ax.set_ylabel("Normalized FFT") - ax.set_title("Please, click on the plot to select cut-off frequency") - plt.tight_layout() - - cid = fig.canvas.mpl_connect("button_press_event", _click_coords) - plt.show() - - cutoff = ix - print(f"Frequency cut-off selected: {ix}") - - # remove high frequency signal and return inverse FFT - FT[freq > cutoff] = 0 - FT[freq < -cutoff] = 0 - - return np.fft.ifft(FT) - - -def _butter_lowpass(cutoff: float, fs: float, order: int = 5): - nyq = 0.5 * fs - normal_cutoff = cutoff / nyq - b, a = butter(order, normal_cutoff, btype="low", analog=False) - return b, a - - -def _butter_lowpass_filter(data: np.ndarray, cutoff: float, fs: float, order: int = 5): - b, a = _butter_lowpass(cutoff, fs, order=order) - y = lfilter(b, a, data) - return y - - -def butter_filtering(data: np.ndarray, dt: float = 1.0, cutoff: Optional[float] = None) -> np.ndarray: - """Filter data using Butterworth filter - - Args: - data: np.ndarray - dt: float = 1.0 - set spacing of data - cutoff: float - set high frequency cut-off (default: None) - """ - - order = 6 - fs = 1 / dt # sample rate - return _butter_lowpass_filter(data, cutoff, fs, order) diff --git a/src/postgkyl/tools/gkeyll_dg_ops.py b/src/postgkyl/tools/gkeyll_dg_ops.py deleted file mode 100644 index 4299855c..00000000 --- a/src/postgkyl/tools/gkeyll_dg_ops.py +++ /dev/null @@ -1,599 +0,0 @@ -""" -Python bindings for Gkeyll DG binary operations via ctypes. - -Usage: - ops = GkeyllDGops("/path/to/gkylsoft/gkeyll") - ops.invert(0, out_gdata, 0, inp_gdata) - ops.multiply(0, out_gdata, 0, lop_gdata, 0, rop_gdata) -""" - -import ctypes -import os - -import numpy as np - -from postgkyl._gkylsoft_path import resolve_gkylsoft_path -from postgkyl.data import GData -import postgkyl.utils.gkeyll_enums as gke -from postgkyl.data.dg import _getnum_nodes -from postgkyl.modalDG.kernels import expand_1d - -# gkyl_elem_type enum ordinal for double (INT=0, FLOAT=1, DOUBLE=2) -_GKYL_DOUBLE = ctypes.c_int(2) - -class GkeyllDGops: - """ - Operations on DG data, returning DG data. - Some of these are implemented in Gkeyll, and we get them from libg0core.so. - - Inputs: - gkylsoft_path: Path to the gkylsoft directory (the one containing gkeyll/lib/libg0core.so). - If None, falls back to the GKYLSOFT env var, ~/.postgkyl/gkylsoft_path, - and the build-time default in postgkyl._gkylsoft_path. - """ - - def __init__(self, gkylsoft_path: str | None = None): - path = resolve_gkylsoft_path(gkylsoft_path) - if path is None: - raise RuntimeError("gkylsoft path not configured. Set the GKYLSOFT environment variable, " - "write the path to ~/.postgkyl/gkylsoft_path, or pass gkylsoft_path= " - "to GkeyllDGops().") - # end - lib_file = os.path.join(path, "gkeyll/lib", "libg0core.so") - if not os.path.isfile(lib_file): - raise FileNotFoundError(f"libg0core.so not found at {lib_file}. " - "Check that the gkylsoft path is correct.") - # end - self._lib = ctypes.CDLL(lib_file) - self._setup_signatures() - - def _setup_signatures(self) -> None: - lib = self._lib - c_vp = ctypes.c_void_p - c_i = ctypes.c_int - c_sz = ctypes.c_size_t - c_d = ctypes.c_double - - # gkyl_array_new_from_buff(type, ncomp, size, buff) -> gkyl_array* - lib.gkyl_array_new_from_buff.argtypes = [c_i, c_sz, c_sz, c_vp] - lib.gkyl_array_new_from_buff.restype = c_vp - - # gkyl_array_release(arr) - lib.gkyl_array_release.argtypes = [c_vp] - lib.gkyl_array_release.restype = None - - # gkyl_cart_modal_serendip_new(ndim, poly_order) -> gkyl_basis* - lib.gkyl_cart_modal_serendip_new.argtypes = [c_i, c_i] - lib.gkyl_cart_modal_serendip_new.restype = c_vp - - # gkyl_cart_modal_gkhybrid_new(cdim, vdim) -> gkyl_basis* - lib.gkyl_cart_modal_gkhybrid_new.argtypes = [c_i, c_i] - lib.gkyl_cart_modal_gkhybrid_new.restype = c_vp - - # gkyl_cart_modal_basis_get_num_basis(*basis) -> int - lib.gkyl_cart_modal_basis_get_num_basis.argtypes = [c_vp] - lib.gkyl_cart_modal_basis_get_num_basis.restype = c_i - - # gkyl_cart_modal_basis_release(basis) - lib.gkyl_cart_modal_basis_release.argtypes = [c_vp] - lib.gkyl_cart_modal_basis_release.restype = None - - # gkyl_rect_grid_new(ndim, *lower, *upper, *cells) -> gkyl_rect_grid* - lib.gkyl_rect_grid_new.argtypes = [c_i, ctypes.POINTER(c_d), - ctypes.POINTER(c_d), ctypes.POINTER(c_i)] - lib.gkyl_rect_grid_new.restype = c_vp - - # gkyl_rect_grid_release(grid) - lib.gkyl_rect_grid_release.argtypes = [c_vp] - lib.gkyl_rect_grid_release.restype = None - - # gkyl_range_new(ndim, *lower, *upper) -> gkyl_range* - lib.gkyl_range_new.argtypes = [c_i, ctypes.POINTER(c_i), ctypes.POINTER(c_i)] - lib.gkyl_range_new.restype = c_vp - - # gkyl_range_release(rng) - lib.gkyl_range_release.argtypes = [c_vp] - lib.gkyl_range_release.restype = None - - # gkyl_dg_mul_op(*basis, c_oop, *out, c_lop, *lop, c_rop, rop*) - lib.gkyl_dg_mul_op.argtypes = [c_vp, c_i, c_vp, c_i, c_vp, c_i, c_vp] - lib.gkyl_dg_mul_op.restype = None - - # gkyl_dg_mul_conf_phase_op_range(*cbasis, *pbasis, *pout, *cop, *pop, *crange, *prange) - lib.gkyl_dg_mul_conf_phase_op_range.argtypes = [c_vp, c_vp, c_vp, c_vp, c_vp, c_vp, c_vp] - lib.gkyl_dg_mul_conf_phase_op_range.restype = None - - # gkyl_dg_inv_op(*basis, c_oop, *out, c_iop, *iop) - lib.gkyl_dg_inv_op.argtypes = [c_vp, c_i, c_vp, c_i, c_vp] - lib.gkyl_dg_inv_op.restype = None - - # gkyl_dg_differentiate_op_local(*basis, dir, diff_order, dx, c_oop, *out, c_iop, inp*) - lib.gkyl_dg_differentiate_op_local.argtypes = [c_vp, c_i, c_i, c_d, c_i, c_vp, c_i, c_vp] - lib.gkyl_dg_differentiate_op_local.restype = None - - # gkyl_dg_eval_at_coord_proj_new(cdim_do, *basis_do, num_eval_dirs, *eval_dirs, use_gpu) - lib.gkyl_dg_eval_at_coord_proj_new.argtypes = [c_i, c_vp, c_i, ctypes.POINTER(c_i), ctypes.c_bool] - lib.gkyl_dg_eval_at_coord_proj_new.restype = c_vp - - # gkyl_dg_eval_at_coord_proj_target_basis(up*, cdim*, ndim*, btype*, poly_order*, num_basis*) - lib.gkyl_dg_eval_at_coord_proj_target_basis.argtypes = [ - c_vp, ctypes.POINTER(c_i), ctypes.POINTER(c_i), ctypes.POINTER(c_i), - ctypes.POINTER(c_i), ctypes.POINTER(c_i), - ] - lib.gkyl_dg_eval_at_coord_proj_target_basis.restype = None - - # gkyl_dg_eval_at_coord_proj_advance(up*, eval_coords*, grid*, pick_lower*, - # known_index*, rng_do*, rng_tar*, fdo*, ftar*) - lib.gkyl_dg_eval_at_coord_proj_advance.argtypes = [ - c_vp, ctypes.POINTER(c_d), c_vp, ctypes.POINTER(ctypes.c_bool), - ctypes.POINTER(c_i), c_vp, c_vp, c_vp, c_vp, - ] - lib.gkyl_dg_eval_at_coord_proj_advance.restype = None - - # gkyl_dg_eval_at_coord_proj_release(up*) - lib.gkyl_dg_eval_at_coord_proj_release.argtypes = [c_vp] - lib.gkyl_dg_eval_at_coord_proj_release.restype = None - - # gkyl_proj_powsqrt_on_basis_new(*basis, num_quad, use_gpu) -> gkyl_proj_powsqrt_on_basis* - lib.gkyl_proj_powsqrt_on_basis_new.argtypes = [c_vp, c_i, ctypes.c_bool] - lib.gkyl_proj_powsqrt_on_basis_new.restype = c_vp - - # gkyl_proj_powsqrt_on_basis_advance(up*, *range, expIn, *fIn, *fOut) - lib.gkyl_proj_powsqrt_on_basis_advance.argtypes = [c_vp, c_vp, c_d, c_vp, c_vp] - lib.gkyl_proj_powsqrt_on_basis_advance.restype = None - - # gkyl_proj_powsqrt_on_basis_release(up*) - lib.gkyl_proj_powsqrt_on_basis_release.argtypes = [c_vp] - lib.gkyl_proj_powsqrt_on_basis_release.restype = None - - # gkyl_array_average_new(*grid, *basis, *basis_avg, *local, *local_avg, - # *local_avg_ext, *weight, *avg_dim, use_gpu) -> gkyl_array_average* - lib.gkyl_array_average_new.argtypes = [c_vp, c_vp, c_vp, c_vp, c_vp, c_vp, c_vp, - ctypes.POINTER(c_i), ctypes.c_bool] - lib.gkyl_array_average_new.restype = c_vp - - # gkyl_array_average_advance(up*, fin*, avgout*) - lib.gkyl_array_average_advance.argtypes = [c_vp, c_vp, c_vp] - lib.gkyl_array_average_advance.restype = None - - # gkyl_array_average_release(up*) - lib.gkyl_array_average_release.argtypes = [c_vp] - lib.gkyl_array_average_release.restype = None - - def _gkyl_array_new_from_gdata(self, gdata): - """ - Wrap a GData's value buffer in a gkyl_array without copying. - - Returns (arr_ptr, values) where values is the numpy array kept alive - to prevent GC while arr_ptr is in use. - """ - values = np.squeeze(gdata.get_values()) - # Ensure C-contiguous float64 layout expected by gkyl kernels - values = np.ascontiguousarray(values, dtype=np.float64) - size = ctypes.c_size_t(int(np.prod(values.shape[:-1]))) - ncomp = ctypes.c_size_t(int(values.shape[-1])) - data_ptr = values.ctypes.data_as(ctypes.c_void_p) - arr_ptr = self._lib.gkyl_array_new_from_buff(_GKYL_DOUBLE, ncomp, size, data_ptr) - return arr_ptr, values - - def _gkyl_basis_new_from_gdata(self, gdata): - """Create a basis from a GData's metadata. Caller must release.""" - ndim = gdata.get_num_dims() - poly_order = int(gdata.ctx["poly_order"]) - basis_type = gdata.ctx["basis_type"] - if basis_type == "gkhybrid": - vdim = 1 if ndim == 2 else 2 - cdim = ndim - vdim - return self._lib.gkyl_cart_modal_gkhybrid_new(ctypes.c_int(cdim), ctypes.c_int(vdim)) - else: - return self._lib.gkyl_cart_modal_serendip_new(ctypes.c_int(ndim), ctypes.c_int(poly_order)) - - def _gkyl_range_new_from_gdata(self, gdata): - """Create a 1-indexed gkyl_range covering all cells of gdata. Caller must release.""" - values = gdata.get_values() - cells = list(values.shape[:-1]) - ndim = len(cells) - c_lo = (ctypes.c_int * ndim)(*([1] * ndim)) - c_up = (ctypes.c_int * ndim)(*cells) - return self._lib.gkyl_range_new(ctypes.c_int(ndim), c_lo, c_up) - - def multiply(self, c_oop: int, oop, c_lop: int, lop, c_rop: int, rop) -> None: - """ - Weak DG multiply: oop[c_oop] = lop[c_lop] * rop[c_rop]. - - Inputs: - c_oop, c_lop, c_rop: Physical component indices (0-based) within each multi-component field. - Use 0 for single-component (scalar) fields. - oop, lop, rop: Output and input operand datasets. Must be pre-allocated. - """ - basis = self._gkyl_basis_new_from_gdata(lop) - arr_oop, _ = self._gkyl_array_new_from_gdata(oop) - arr_lop, _ = self._gkyl_array_new_from_gdata(lop) - arr_rop, _ = self._gkyl_array_new_from_gdata(rop) - try: - self._lib.gkyl_dg_mul_op(basis, - ctypes.c_int(c_oop), arr_oop, - ctypes.c_int(c_lop), arr_lop, - ctypes.c_int(c_rop), arr_rop,) - finally: - self._lib.gkyl_cart_modal_basis_release(basis) - self._lib.gkyl_array_release(arr_oop) - self._lib.gkyl_array_release(arr_lop) - self._lib.gkyl_array_release(arr_rop) - - def multiply_conf_phase(self, pout, cop, pop) -> None: - """ - Weak DG conf-phase multiply: pout = cop * pop on all cells. - - cop is a conf-space field and pop/pout are phase-space fields. - Ranges are constructed automatically from the shape of each dataset. - - Inputs: - pout: Output phase-space dataset. Must be pre-allocated. - cop: Conf-space operand dataset. - pop: Phase-space operand dataset. - """ - cbasis = self._gkyl_basis_new_from_gdata(cop) - pbasis = self._gkyl_basis_new_from_gdata(pop) - arr_pout, _ = self._gkyl_array_new_from_gdata(pout) - arr_cop, _ = self._gkyl_array_new_from_gdata(cop) - arr_pop, _ = self._gkyl_array_new_from_gdata(pop) - crange = self._gkyl_range_new_from_gdata(cop) - prange = self._gkyl_range_new_from_gdata(pop) - try: - self._lib.gkyl_dg_mul_conf_phase_op_range( - cbasis, pbasis, arr_pout, arr_cop, arr_pop, crange, prange) - finally: - self._lib.gkyl_cart_modal_basis_release(cbasis) - self._lib.gkyl_cart_modal_basis_release(pbasis) - self._lib.gkyl_array_release(arr_pout) - self._lib.gkyl_array_release(arr_cop) - self._lib.gkyl_array_release(arr_pop) - self._lib.gkyl_range_release(crange) - self._lib.gkyl_range_release(prange) - - def differentiate(self, dir: int, diff_order: int, dx: float, c_oop: int, oop, c_iop: int, iop) -> None: - """ - Local DG differentiation: oop[c_oop] = d^diff_order/dx_dir^diff_order iop[c_iop]. - - Differentiates the DG expansion in each cell independently (no inter-cell stencil). - - Inputs: - dir: Direction of differentiation (0-based). - diff_order: Order of the derivative (1 or 2). - dx: Cell length in the direction of differentiation. - c_oop, c_iop: Physical component indices (0-based). - oop, iop: Output and input datasets. oop must be allocated. - """ - basis = self._gkyl_basis_new_from_gdata(iop) - arr_oop, _ = self._gkyl_array_new_from_gdata(oop) - arr_iop, _ = self._gkyl_array_new_from_gdata(iop) - try: - self._lib.gkyl_dg_differentiate_op_local(basis, - ctypes.c_int(dir), ctypes.c_int(diff_order), ctypes.c_double(dx), - ctypes.c_int(c_oop), arr_oop, - ctypes.c_int(c_iop), arr_iop,) - finally: - self._lib.gkyl_cart_modal_basis_release(basis) - self._lib.gkyl_array_release(arr_oop) - self._lib.gkyl_array_release(arr_iop) - - def eval_at_coord_proj(self, eval_dirs: list, eval_coords: list, gdata, - comp_grid: bool = False) -> GData: - """ - Evaluate a DG field at physical coordinates in eval_dirs and project onto - the lower-dimensional target basis. - - Inputs: - eval_dirs: Sorted list of 0-based direction indices to eliminate. - eval_coords: Physical coordinates, one per entry in eval_dirs. - gdata: Donor DG dataset (must have poly_order in ctx). - comp_grid: Passed to the output GData constructor. - - Returns: - GData with the projected field. The surviving grid dimensions, cells, - lower, upper, and num_comps in ctx are set correctly for the target. - """ - ndim = gdata.get_num_dims() - vals = gdata.get_values() - poly_order = int(gdata.ctx["poly_order"]) - - basis_type = gdata.ctx["basis_type"] - grid_type = gdata.ctx["grid_type"] - - ggrid = gdata.get_grid() - grid_edges = [np.copy(ggrid[d]) for d in range(ndim)] - if basis_type == "gkhybrid" and grid_type == "c2p_vel": - # Grid has DG coefficients of v-space mapping along v-dims. Evaluate at cell boundaries. - # MF 2026/06/28: I think this should happen outside of this function, - # but we do it here for now to avoid modifying other code. - poly_order_vmap = 1 - num_cdim = gdata.ctx["num_cdim"] - num_vdim = gdata.ctx["num_vdim"] - num_basis_1v = int(_getnum_nodes(1, 1, "serendipity")) # 1D p1 basis for single v dimension. - nodes = [-1.0, 1.0] - for d in range(num_vdim): - q = grid_edges[num_cdim+d] - grid_edges_1v = np.zeros(np.size(q,0)+1) - for i, vmap_c in enumerate(q): - grid_edges_1v[i] = expand_1d[int(poly_order_vmap - 1)](vmap_c, nodes[0]) - # end - # Append upper boundary surface. - grid_edges_1v[-1] = expand_1d[int(poly_order_vmap - 1)](q[-1], nodes[1]) - - grid_edges[num_cdim+d] = grid_edges_1v - # end - # end - - cells = [len(grid_edges[d]) - 1 for d in range(ndim)] - lower = [float(grid_edges[d][0]) for d in range(ndim)] - upper = [float(grid_edges[d][-1]) for d in range(ndim)] - - num_eval = len(eval_dirs) - keep_dirs = [d for d in range(ndim) if d not in eval_dirs] - ndim_tar = len(keep_dirs) - cells_tar = [cells[d] for d in keep_dirs] if num_eval 0: - c_rng_lo_tar = (ctypes.c_int * ndim_tar)(*([1] * ndim_tar)) - c_rng_up_tar = (ctypes.c_int * ndim_tar)(*cells_tar) - rng_tar_ptr = self._lib.gkyl_range_new(ctypes.c_int(ndim_tar), c_rng_lo_tar, c_rng_up_tar) - tar_grid = [ggrid[d] for d in keep_dirs] # Use original grid to keep mapping if c2p_vel. - else: - c_one = (ctypes.c_int * 1)(1) - rng_tar_ptr = self._lib.gkyl_range_new(ctypes.c_int(1), c_one, c_one) - tar_grid = [np.array([eval_coords[d]]) for d in range(num_eval)] - - # Donor array. - arr_do, values = self._gkyl_array_new_from_gdata(gdata) - ncomp_raw = int(values.shape[-1]) - - # Target buffer. - num_phys_comps = ncomp_raw // num_basis_do - ncomp_tar = num_phys_comps * num_basis_tar - size_tar = int(np.prod(cells_tar)) - tar_shape = (*cells_tar, ncomp_tar) - tar_buf = np.zeros(tar_shape, dtype=np.float64) - arr_tar = self._lib.gkyl_array_new_from_buff(_GKYL_DOUBLE, ctypes.c_size_t(ncomp_tar), - ctypes.c_size_t(size_tar), tar_buf.ctypes.data_as(ctypes.c_void_p), ) - - c_eval_coords = (ctypes.c_double * num_eval)(*eval_coords) - c_pick_lower = (ctypes.c_bool * num_eval)(*([False] * num_eval)) - c_known_idx = (ctypes.c_int * ndim)(*([-1] * ndim)) - try: - self._lib.gkyl_dg_eval_at_coord_proj_advance(updater, c_eval_coords, grid_ptr, - c_pick_lower, c_known_idx, rng_do_ptr, rng_tar_ptr, arr_do, arr_tar,) - finally: - self._lib.gkyl_dg_eval_at_coord_proj_release(updater) - self._lib.gkyl_cart_modal_basis_release(basis_do_ptr) - self._lib.gkyl_array_release(arr_do) - self._lib.gkyl_array_release(arr_tar) - self._lib.gkyl_rect_grid_release(grid_ptr) - self._lib.gkyl_range_release(rng_do_ptr) - self._lib.gkyl_range_release(rng_tar_ptr) - - out = GData(ctx=gdata.ctx, comp_grid=comp_grid) - out.push(tar_grid, tar_buf) - - # Re-set the basis in the context in case it changed. - out.ctx["basis_type"] = gke.basis_type_gkyl_to_pgkyl(int(_btype_tar.value)) - out.ctx["poly_order"] = int(_poly_order_tar.value) - out.ctx["num_cdim"] = int(_cdim_tar.value) - out.ctx["num_vdim"] = int(_ndim_tar.value - _cdim_tar.value) - - return out - - def average(self, avg_dirs: list, gdata, weight=None, comp_grid: bool = False) -> GData: - """ - Average a DG field over the directions in avg_dirs (gkyl_array_average). - - Returns a GData over the surviving dimensions. With a weight GData (same - dims/basis as gdata) the weighted average is computed instead. Serendipity - basis, poly_order <= 2 only. - """ - basis_type = gdata.ctx["basis_type"] - if basis_type.lower() != "serendipity": - raise ValueError(f"average only supports the serendipity basis, got '{basis_type}'. " - "gkyl_array_average provides serendipity kernels only.") - - ndim = gdata.get_num_dims() - poly_order = int(gdata.ctx["poly_order"]) - if poly_order > 2: - raise ValueError(f"average only supports poly_order <= 2, got {poly_order}.") - - if weight is not None: - w_basis_type = weight.ctx["basis_type"] - if w_basis_type.lower() != "serendipity": - raise ValueError(f"weight must use the serendipity basis, got '{w_basis_type}'.") - if weight.get_num_dims() != ndim: - raise ValueError(f"weight has {weight.get_num_dims()} dims but the field has {ndim}; " - "they must match.") - if int(weight.ctx["poly_order"]) != poly_order: - raise ValueError(f"weight poly_order {int(weight.ctx['poly_order'])} != field " - f"poly_order {poly_order}.") - - avg_dirs = sorted(set(avg_dirs)) - if not avg_dirs or avg_dirs[0] < 0 or avg_dirs[-1] >= ndim: - raise ValueError(f"average dirs {avg_dirs} out of range for a {ndim}D field.") - keep_dirs = [d for d in range(ndim) if d not in avg_dirs] - ndim_tar = len(keep_dirs) - - ggrid = gdata.get_grid() - grid_edges = [np.copy(ggrid[d]) for d in range(ndim)] - cells = [len(grid_edges[d]) - 1 for d in range(ndim)] - lower = [float(grid_edges[d][0]) for d in range(ndim)] - upper = [float(grid_edges[d][-1]) for d in range(ndim)] - - # For a full average (no surviving dims), Gkeyll keeps a 1D, single-cell - # target following the same convention as eval_at_coord_proj. - ndim_red = ndim_tar if ndim_tar > 0 else 1 - cells_tar = [cells[d] for d in keep_dirs] if ndim_tar > 0 else [1] - - # Donor grid. - c_lower = (ctypes.c_double * ndim)(*lower) - c_upper = (ctypes.c_double * ndim)(*upper) - c_cells = (ctypes.c_int * ndim)(*cells) - grid_ptr = self._lib.gkyl_rect_grid_new(ctypes.c_int(ndim), c_lower, c_upper, c_cells) - - # Donor (full) range, 1-indexed. - c_rng_lo = (ctypes.c_int * ndim)(*([1] * ndim)) - c_rng_up = (ctypes.c_int * ndim)(*cells) - rng_ptr = self._lib.gkyl_range_new(ctypes.c_int(ndim), c_rng_lo, c_rng_up) - - # Target (reduced) range, 1-indexed. - c_rng_lo_tar = (ctypes.c_int * ndim_red)(*([1] * ndim_red)) - c_rng_up_tar = (ctypes.c_int * ndim_red)(*cells_tar) - rng_tar_ptr = self._lib.gkyl_range_new(ctypes.c_int(ndim_red), c_rng_lo_tar, c_rng_up_tar) - - # Full (donor) and reduced (target) serendipity bases. - basis_do = self._gkyl_basis_new_from_gdata(gdata) - basis_avg = self._lib.gkyl_cart_modal_serendip_new(ctypes.c_int(ndim_red), - ctypes.c_int(poly_order)) - num_basis_do = int(self._lib.gkyl_cart_modal_basis_get_num_basis(basis_do)) - num_basis_tar = int(self._lib.gkyl_cart_modal_basis_get_num_basis(basis_avg)) - - # Donor array. - arr_do, values = self._gkyl_array_new_from_gdata(gdata) - ncomp_raw = int(values.shape[-1]) - - # Optional weight array (spans the full donor range/basis). Keep _w_values - # alive so its numpy buffer is not collected while the kernel runs. - arr_w, _w_values = (None, None) - if weight is not None: - arr_w, _w_values = self._gkyl_array_new_from_gdata(weight) - - # Target buffer. - num_phys_comps = ncomp_raw // num_basis_do - ncomp_tar = num_phys_comps * num_basis_tar - size_tar = int(np.prod(cells_tar)) - tar_shape = (*cells_tar, ncomp_tar) - tar_buf = np.zeros(tar_shape, dtype=np.float64) - arr_tar = self._lib.gkyl_array_new_from_buff(_GKYL_DOUBLE, ctypes.c_size_t(ncomp_tar), - ctypes.c_size_t(size_tar), tar_buf.ctypes.data_as(ctypes.c_void_p), ) - - # avg_dim flags (1 = averaged) over the full dimensionality. - avg_flags = [1 if d in avg_dirs else 0 for d in range(ndim)] - c_avg_dim = (ctypes.c_int * ndim)(*avg_flags) - - # rng_tar_ptr doubles as local_avg_ext (only read to size the integrated weight). - updater = self._lib.gkyl_array_average_new(grid_ptr, basis_do, basis_avg, - rng_ptr, rng_tar_ptr, rng_tar_ptr, arr_w, c_avg_dim, ctypes.c_bool(False)) - try: - self._lib.gkyl_array_average_advance(updater, arr_do, arr_tar) - finally: - self._lib.gkyl_array_average_release(updater) - self._lib.gkyl_cart_modal_basis_release(basis_do) - self._lib.gkyl_cart_modal_basis_release(basis_avg) - self._lib.gkyl_array_release(arr_do) - self._lib.gkyl_array_release(arr_tar) - if arr_w is not None: - self._lib.gkyl_array_release(arr_w) - self._lib.gkyl_rect_grid_release(grid_ptr) - self._lib.gkyl_range_release(rng_ptr) - self._lib.gkyl_range_release(rng_tar_ptr) - - tar_grid = [ggrid[d] for d in keep_dirs] if ndim_tar > 0 else [np.array([0.0, 1.0])] - - out = GData(ctx=gdata.ctx, comp_grid=comp_grid) - out.push(tar_grid, tar_buf) - - out.ctx["basis_type"] = "serendipity" - out.ctx["poly_order"] = poly_order - out.ctx["num_cdim"] = ndim_tar - out.ctx["num_vdim"] = 0 - - return out - - def invert(self, c_oop: int, oop, c_iop: int, iop) -> None: - """ - Weak DG invert: oop[c_oop] = 1 / iop[c_iop]. - - Only supported for serendipity basis at p=1 (gkeyll limitation). - - Inputs: - c_oop, c_iop: Physical component indices (0-based). - oop, iop: Output and input datasets. oop be allocated. - """ - basis = self._gkyl_basis_new_from_gdata(iop) - arr_oop, _ = self._gkyl_array_new_from_gdata(oop) - arr_iop, _ = self._gkyl_array_new_from_gdata(iop) - try: - self._lib.gkyl_dg_inv_op(basis, - ctypes.c_int(c_oop), arr_oop, - ctypes.c_int(c_iop), arr_iop,) - finally: - self._lib.gkyl_cart_modal_basis_release(basis) - self._lib.gkyl_array_release(arr_oop) - self._lib.gkyl_array_release(arr_iop) - - def powsqrt(self, oop, iop, exponent: float, num_quad: int | None = None) -> None: - """ - Weak DG power of a square root: oop = pow(sqrt(iop), exponent), projected - onto the basis by Gauss-Legendre quadrature (gkyl_proj_powsqrt_on_basis). - - Inputs: - oop, iop: Output and input datasets. oop must be pre-allocated. - exponent: Exponent applied to sqrt(iop). - num_quad: Quadrature nodes per direction. Defaults to poly_order+1, - matching the gyrokinetic app's own use of this updater. - """ - basis = self._gkyl_basis_new_from_gdata(iop) - try: - num_basis = int(self._lib.gkyl_cart_modal_basis_get_num_basis(basis)) - for name, gdata in (("iop", iop), ("oop", oop)): - num_comps = int(np.squeeze(gdata.get_values()).shape[-1]) - if num_comps != num_basis: - raise ValueError( - f"powsqrt: '{name}' has {num_comps} coefficients per cell but the basis has " - f"{num_basis}; this operation only takes single-component (scalar) fields.") - # end - # end - - if num_quad is None: - num_quad = int(iop.ctx["poly_order"]) + 1 - # end - - up = self._lib.gkyl_proj_powsqrt_on_basis_new( - basis, ctypes.c_int(num_quad), ctypes.c_bool(False)) - arr_oop, _ = self._gkyl_array_new_from_gdata(oop) - arr_iop, _ = self._gkyl_array_new_from_gdata(iop) - rng = self._gkyl_range_new_from_gdata(iop) - try: - self._lib.gkyl_proj_powsqrt_on_basis_advance( - up, rng, ctypes.c_double(exponent), arr_iop, arr_oop) - finally: - self._lib.gkyl_proj_powsqrt_on_basis_release(up) - self._lib.gkyl_array_release(arr_oop) - self._lib.gkyl_array_release(arr_iop) - self._lib.gkyl_range_release(rng) - finally: - self._lib.gkyl_cart_modal_basis_release(basis) - diff --git a/src/postgkyl/tools/growth.py b/src/postgkyl/tools/growth.py deleted file mode 100644 index f5eb1cf7..00000000 --- a/src/postgkyl/tools/growth.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Postgkyl module for fitting growth rates.""" - -import numpy as np -import scipy.optimize as opt -import sys -from typing import Callable, Tuple - - -def exp2(x: float, a: float, b: float) -> float: - """Define custom exponential a*exp(2b*x) - - Args: - x: float - independent variable - a: float - scaling parameter - b: float - growth rate - - Notes: - Energy (quantity^2) is often used for the growth-rate study, - therefore the factor 2 - """ - return a*np.exp(2*b*x) - - -def fit_growth(x: np.ndarray, y: np.ndarray, function: Callable = exp2, - min_N: int | None = None, p0: tuple = (1, 1)) -> Tuple[tuple, float, int]: - """Fit function to continuously increasing region of data - - Parameters: - x: NumPy array - independet variable - y: NumPy array - dependent variable - min_N: int - minimal number of fitted points - function: callable = exp2 - function to fit - p0: tuple = (1, 1) - initial guess - - Notes: - The best is determined based on the coeficient of determination, - R^2 https://en.wikipedia.org/wiki/Coefficient_of_determination - """ - best_R2 = 0.0 - if min_N is None: - min_N = int(len(x)/10) - max_N = len(x) - best_N = min_N - best_params = p0 - - max_x = x[-1] - - print(f"fit_growth: fitting region {min_N:d} -> {max_N:d}") - for n in np.linspace(min_N, max_N - 1, max_N - min_N): - n = int(n) - xn = x[0:n]/max_x # continuously increasing fitting region - yn = y[0:n] - try: - params, _ = opt.curve_fit(function, xn, yn, best_params) - residual = yn - function(xn, *params) - ss_res = np.sum(residual**2) - ss_tot = np.sum((yn - np.mean(yn))**2) - R2 = 1 - ss_res/ss_tot - if R2 > best_R2: - best_R2 = R2 - best_params = params - best_N = n - # end - percent = float(n - min_N) / (max_N - min_N)*100 - progress = "[" + int(percent / 10) * "=" + (10 - int(percent / 10)) * " " + "]" - sys.stdout.write( - f"\rgamma = {best_params[1] / max_x:+.5e} (current {params[1] / max_x:+.3e} R^2={R2:.3e}) {percent:6.2f}% done {progress}") - sys.stdout.flush() - except RuntimeError: - print(f"fit_growth: curve_fit failed for N = {n:d}") - # end - # end - best_params[1] = best_params[1]/max_x - print(f"\ngamma = {best_params[1]:+.5e}") - return best_params, best_R2, best_N diff --git a/src/postgkyl/tools/init_polar.py b/src/postgkyl/tools/init_polar.py deleted file mode 100644 index 5f54af86..00000000 --- a/src/postgkyl/tools/init_polar.py +++ /dev/null @@ -1,64 +0,0 @@ -import numpy as np - - -def init_polar(nkx, nky, nkz, kx, ky, kz, nkpolar): - # if 2D, nkz and kz = 0 - - if nkpolar == 0: - akp = [] - nbin = 0 - polar_index = [] - akplim = [] - elif nkz == 0: - nbin = np.zeros(nkpolar) # Number of kx,ky in each polar bins - polar_index = np.zeros((nkx, nky), dtype=int) # Polar index to simplify binning - if nkx == 1 & nky == 1: - dkp = 0 - elif nkx == 1: - dkp = ky[1] - elif nky == 1: - dkp = kx[1] - else: - dkp = max(kx[1], ky[1]) - akp = (np.linspace(1, nkpolar, nkpolar)) * dkp # Kperp grid - akplim = dkp / 2 + (np.linspace(0, nkpolar, nkpolar + 1))*dkp # Bin limits - # Re-written to avoid loops. Necessary for large grids. - [kxg, kyg] = np.meshgrid( - ky, kx - ) # Deal with meshgrid weirdness (so do not have to transpose) - kp = np.sqrt(kxg**2 + kyg**2) - pn = np.where(kp >= akplim[nkpolar]) - polar_index[pn[0], pn[1]] = nkpolar - 1 - nbin[nkpolar - 1] = nbin[nkpolar - 1] + len(pn[0]) - for ik in range(0, nkpolar): - pn = np.where((kp < akplim[ik + 1]) & (kp >= akplim[ik])) - polar_index[pn[0], pn[1]] = ik - nbin[ik] = nbin[ik] + len(pn[0]) - else: - # 3D data - nbin = np.zeros(nkpolar) - polar_index = np.zeros((nkx, nky, nkz), dtype=int) - if nkx == 1 & nky == 1 & nkz == 1: - dkp = 0 - elif nkx == 1: - dkp = max(ky[1], kz[1]) - elif nky == 1: - dkp = max(kx[1], kz[1]) - elif nkz == 1: - dkp = max(kx[1], ky[1]) - else: - dkp = max(kx[1], ky[1], kz[1]) - akp = (np.linspace(1, nkpolar, nkpolar)) * dkp # kperp grid - akplim = dkp / 2 + (np.linspace(0, nkpolar, nkpolar + 1)) * dkp # bin limits - # Re-written to avoid loops - [kxg, kyg, kzg] = np.meshgrid(ky, kx, kz) - kp = np.sqrt(kxg**2 + kyg**2 + kzg**2) - pn = np.where(kp >= akplim[nkpolar]) - polar_index[pn[0], pn[1], pn[2]] = nkpolar - 1 - nbin[nkpolar - 1] = nbin[nkpolar - 1] + len(pn[0]) - for ik in range(0, nkpolar): - pn = np.where((kp < akplim[ik + 1]) & (kp >= akplim[ik])) - polar_index[pn[0], pn[1], pn[2]] = ik - nbin[ik] = nbin[ik] + len(pn[0]) - - return akp, nbin, polar_index, akplim diff --git a/src/postgkyl/tools/laguerre_compose.py b/src/postgkyl/tools/laguerre_compose.py deleted file mode 100644 index 0093f912..00000000 --- a/src/postgkyl/tools/laguerre_compose.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Postgkyl module for combining the two laguerre components F0 and F1. - -Within Gkeyll, this is mostly use for working with the PKPM data. -""" -from __future__ import annotations - -import numpy as np -from typing import Tuple, TYPE_CHECKING - -from postgkyl.utils import input_parser -if TYPE_CHECKING: - from postgkyl import GData -# end - - -def laguerre_compose(in_f: GData | Tuple[list, np.ndarray], - in_T_m: GData | Tuple[list, np.ndarray], - out_f: GData | None = None) -> Tuple[list, np.ndarray]: - """Compose PKPM expansion coefficients into a single f. - - Compose the full distribution function f(x, v_par, v_perp) out of the - Laguerre expansion coefficients F0(x, v_par), F1(x, v_par) and the - PKPM moments to calculate the T(x) over m. - - Jimmy Juno's slides: https://drive.google.com/file/d/1548tLF9o7vyW3bkrsq6FvAMV-8XJvKtY/view - - Args: - in_f: GData or NumPy array - 2-component Laguerre expansion coefficients. - in_T_m: GData or NumPy array - PKPM T over m. - out_f: GData = None - (Optional) GData to store output. - - Returns: - A tuple of grid (which is itself a tuple of nupy arrays for each dimension) and a - NumPy array with values. - """ - in_f_grid, in_f_values = input_parser(in_f) - _, in_T_m_values = input_parser(in_T_m) - - x, vpar = in_f_grid[0], in_f_grid[1] - vperp = np.copy(vpar) - - x_cc = (x[:-1] + x[1:])/2 - vpar_cc = (vpar[:-1] + vpar[1:])/2 - vperp_cc = (vpar[:-1] + vpar[1:])/2 - - _, _, vperp_3D = np.meshgrid(x_cc, vpar_cc, vperp_cc, indexing="ij") - - F0 = in_f_values[..., 0] - G = in_f_values[..., 1] - T_m = in_T_m_values[..., 0] - - F1 = F0 - (G.transpose()/T_m).transpose() - - # Ading the np.newaxis allows the subsequent np.multiply (called when - # doing * on numpy arrays) to work. The arrays need to have the same - # number of axis, e.g., one can not multiply (3, 3) and (3,) arrays - # but can multiply (3, 3) with (3, 1) or (1, 3). - F0, F1 = F0[..., np.newaxis], F1[..., np.newaxis] - T_m = T_m[..., np.newaxis, np.newaxis] - - # Hardcoded for l=0, n=0,1 in - # https://drive.google.com/file/d/1548tLF9o7vyW3bkrsq6FvAMV-8XJvKtY/view - f = (F0 + F1*(1 - vperp_3D**2/2/T_m))/(2*np.pi*T_m) * np.exp(-(vperp_3D**2)/2/T_m) - - f = f[..., np.newaxis] # Adding the component index - - if out_f: - out_f.push([x, vpar, vperp], f) - # end - return [x, vpar, vperp], f diff --git a/src/postgkyl/tools/mag_sq.py b/src/postgkyl/tools/mag_sq.py deleted file mode 100644 index f8c32896..00000000 --- a/src/postgkyl/tools/mag_sq.py +++ /dev/null @@ -1,42 +0,0 @@ -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -from postgkyl.utils import input_parser -if TYPE_CHECKING: - from postgkyl import GData -# end - - -def mag_sq(dat: GData | Tuple[list, np.ndarray], coords: str = "0:3", - output: GData | None = None) -> Tuple[list, np.ndarray]: - """Function to compute the magnitude squared of an array - - Parameters: - data - input GData data structure - coords - specific coordinates to compute magnitude squared of by default assume a three - component field and that you want the magnitude squared of the those three - components - - Notes: - Assumes that the number of components is the last dimension. - - """ - in_grid, in_values = input_parser(dat) - - # Because coords is an input string, need to split and parse it to get the right - # coordinates. - s = coords.split(":") - values = in_values[..., slice(int(s[0]), int(s[1]))] - # Output is a scalar, so dimensionality should not include number of components. - out = np.zeros(values[..., 0].shape) - out = np.sum(values*values, axis=-1) - out = out[..., np.newaxis] - - if output: - output.push(in_grid, out) - # end - return in_grid, out diff --git a/src/postgkyl/tools/params.py b/src/postgkyl/tools/params.py deleted file mode 100644 index 6919fa9c..00000000 --- a/src/postgkyl/tools/params.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Postgkyl module for plasma related parameters.""" - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -from postgkyl.tools.mag_sq import mag_sq -from postgkyl.tools.prim_vars import get_density, get_temp, get_mhd_temp -from postgkyl.utils import input_parser - -if TYPE_CHECKING: - from postgkyl import GData -# end - - -def get_magB(field: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - field_grid, field_values = input_parser(field) - b_values = field_values[..., 3:6] - _, mag_B_sq = mag_sq((field_grid, b_values)) - out_values = np.sqrt(mag_B_sq) - - return field_grid, out_values - - -def get_vt(species: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3.0, - num_moms : int | None = None, mass: float = 1.0, mu_0: float = 1.0, - sqrt2: bool = True, mhd: bool = False) -> Tuple[list, np.ndarray]: - m = species.ctx["mass"] if species.ctx["mass"] else mass - - if mhd: - out_grid, temp = get_mhd_temp(species, gas_gamma=gas_gamma, mu_0=mu_0) - else: - out_grid, temp = get_temp(species, gas_gamma=gas_gamma, num_moms=num_moms) - # end - out_values = np.sqrt(temp/m) - if sqrt2: - out_values *= np.sqrt(2.0) - - return out_grid, out_values - - -def get_vA(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - mu_0: float = 1.0) -> Tuple[list, np.ndarray]: - mu = field.ctx["mu_0"] if field.ctx["mu_0"] else mu_0 - - _, magB = get_magB(field) - # Fluid data already has mass factor in density - out_grid, rho = get_density(species) - out_values = magB/np.sqrt(mu*rho) - - return out_grid, out_values - - -def get_omegaC(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - mass: float = 1.0, charge: float = 1.0) -> Tuple[list, np.ndarray]: - m = species.ctx["mass"] if species.ctx["mass"] else mass - q = species.ctx["charge"] if species.ctx["charge"] else charge - - out_grid, magB = get_magB(field) - out_values = abs(q)*magB/m - - return out_grid, out_values - - -def get_omegaP(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - mass: float = 1.0, charge: float = 1.0, epsilon_0: float = 1.0) -> Tuple[list, np.ndarray]: - m = species.ctx["mass"] if species.ctx["mass"] else mass - q = species.ctx["charge"] if species.ctx["charge"] else charge - epsilon = field.ctx["epsilon_0"] if field.ctx["epsilon_0"] else epsilon_0 - - # Fluid data already has mass factor in density - out_grid, rho = get_density(species) - qbym2 = q**2/m**2 - out_values = np.sqrt(qbym2*rho/epsilon) - - return out_grid, out_values - - -def get_d(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - mass: float = 1.0, charge: float = 1.0, epsilon_0: float = 1.0, - mu_0 : float = 1.0) -> Tuple[list, np.ndarray]: - epsilon = field.ctx["epsilon_0"] if field.ctx["epsilon_0"] else epsilon_0 - mu = field.ctx["mu_0"] if field.ctx["mu_0"] else mu_0 - - out_grid, omegaP = get_omegaP(species=species, field=field, mass=mass, charge=charge, - epsilon_0=epsilon_0) - light_speed = 1.0/np.sqrt(epsilon*mu) - out_values = light_speed/omegaP - - return out_grid, out_values - - -def get_lambdaD(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - gas_gamma: float = 5.0/3.0, num_moms: int | None = None, - mass: float = 1.0, charge: float = 1.0, epsilon_0: float = 1.0, - mu_0 : float = 1.0, sqrt2: float = True) -> Tuple[list, np.ndarray]: - _, omegaP = get_omegaP(species=species, field=field, mass=mass, charge=charge, - epsilon_0=epsilon_0) - out_grid, vt = get_vt(species=species, gas_gamma=gas_gamma, num_moms=num_moms, - mass=mass, mu_0=mu_0, sqrt2=sqrt2) - out_values = vt / omegaP - if sqrt2: - out_values /= np.sqrt(2.0) - # end - - return out_grid, out_values - - -def get_rho(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - gas_gamma: float = 5.0/3.0, num_moms: int | None = None, - mass: float = 1.0, charge: float = 1.0, epsilon_0: float = 1.0, - mu_0 : float = 1.0, sqrt2: float = True) -> Tuple[list, np.ndarray]: - - _, omegaC = get_omegaC(species=species, field=field, mass=mass, charge=charge) - out_grid, vt = get_vt(species=species, gas_gamma=gas_gamma, num_moms=num_moms, - mass=mass, mu_0=mu_0, sqrt2=sqrt2) - - out_values = vt/omegaC - if not sqrt2: - out_values *= np.sqrt(2.0) - # end - - return out_grid, out_values - - -def get_beta(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - gas_gamma: float = 5.0/3.0, num_moms: int | None = None, - mass: float = 1.0, mu_0 : float = 1.0, sqrt2: float = True) -> Tuple[list, np.ndarray]: - _, v_A = get_vA(species=species, field=field, mu_0=mu_0) - out_grid, vt = get_vt(species=species, gas_gamma=gas_gamma, num_moms=num_moms, - mass=mass, mu_0=mu_0, sqrt2=sqrt2) - out_values = vt**2 / v_A**2 - if not sqrt2: - out_values *= 2.0 - - return out_grid, out_values diff --git a/src/postgkyl/tools/parrotate.py b/src/postgkyl/tools/parrotate.py deleted file mode 100644 index 6d8af9a9..00000000 --- a/src/postgkyl/tools/parrotate.py +++ /dev/null @@ -1,54 +0,0 @@ - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -if TYPE_CHECKING: - from postgkyl import GData -#end - - -def parrotate(data: GData, rotator: GData, rotate_coords: str = "0:3", - overwrite=False, stack=False) -> Tuple[list, np.ndarray]: - """Function to rotate input array into coordinate system parallel to rotator array - For two arrays u and v, where v is the rotator, operation is (u dot v_hat) v_hat. - - Parameters: - data -- input GData object being rotated - rotator -- GData object used for the rotation - rotate_coords -- optional input to specify a different set of coordinates in the rotator array used - for the rotation (e.g., if rotating to the local magnetic field of a finite volume simulation, rotate_coords='3:6') - - Notes: - Assumes three component fields, and that the number of components is the last dimension. - For a three-component field, the output is a new vector - whose components are (u_{v_x}, u_{v_y}, u_{v_z}), i.e., - the x, y, and z components of the vector u parallel to v. - """ - if stack: - overwrite = stack - print("Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'") - # end - grid = data.get_grid() - values = data.get_values() - # Because rotate_coords is an input string, need to split and parse it to get the right coordinates - s = rotate_coords.split(":") - valuesrot = rotator.get_values()[..., slice(int(s[0]), int(s[1]))] - - outrot = np.zeros(values.shape) - # Assumes three component fields and that the number of components is the last dimension - try: - outrot[..., 0] = np.sum(values*valuesrot, axis=-1)/(np.sum(valuesrot*valuesrot, axis=-1))*valuesrot[..., 0] - outrot[..., 1] = np.sum(values*valuesrot, axis=-1)/(np.sum(valuesrot*valuesrot, axis=-1))*valuesrot[..., 1] - outrot[..., 2] = np.sum(values*valuesrot, axis=-1)/(np.sum(valuesrot*valuesrot, axis=-1))*valuesrot[..., 2] - except IndexError: - print( - f"parrotate: rotation failed due to different numbers of components, data numComponets = '{values.shape[-1]:d}', rotator numComponents = '{rotator.shape[-1]:d}'" - ) - quit() - # end - if overwrite: - data.push(grid, outrot) - - return grid, outrot diff --git a/src/postgkyl/tools/perprotate.py b/src/postgkyl/tools/perprotate.py deleted file mode 100644 index 160a9c8c..00000000 --- a/src/postgkyl/tools/perprotate.py +++ /dev/null @@ -1,44 +0,0 @@ - -from __future__ import annotations - -import numpy as np -from typing import Tuple, TYPE_CHECKING - -from postgkyl.tools.parrotate import parrotate -if TYPE_CHECKING: - from postgkyl import GData -#end - - - -def perprotate(data: GData, rotator: GData, rotate_coords: str = "0:3", - overwrite=False, stack=False) -> Tuple[list, np.ndarray]: - """Function to rotate input array into coordinate system perpendicular to rotator array - For two arrays u and v, where v is the rotator, operation is u - (u dot v_hat) v_hat. - Uses the diagnostic parrotate.py to compute (u dot v_hat) v_hat. - - Parameters: - data -- input GData object being rotated - rotator -- GData object used for the rotation - rotate_coords -- optional input to specify a different set of coordinates in the rotator array used - for the rotation (e.g., if rotating to the local magnetic field of a finite volume simulation, rotate_coords='3:6') - - Notes: - Assumes three component fields, and that the number of components is the last dimension. - """ - if stack: - overwrite = stack - print( - "Deprecation warning: The 'stack' parameter is going to be replaced with 'overwrite'" - ) - # end - grid = data.get_grid() - values = data.get_values() - - outrot = np.zeros_like(values) - outrot = values - parrotate(data, rotator, rotate_coords) - if overwrite: - data.push(grid, outrot) - #end - - return grid, outrot diff --git a/src/postgkyl/tools/polar_isotropic.py b/src/postgkyl/tools/polar_isotropic.py deleted file mode 100644 index 7389c593..00000000 --- a/src/postgkyl/tools/polar_isotropic.py +++ /dev/null @@ -1,19 +0,0 @@ -import numpy as np - - -def polar_isotropic(nkpolar, nkx, nky, nkz, polar_index, nbin, fft_matrix, kx, ky, kz): - # if 2D, then nkz = kz = 0 - - fft_isok = np.zeros(nkpolar) - if nkz == 0: - for i in range(0, nkx): - for j in range(0, nky): - fft_isok[polar_index[i, j]] = fft_isok[polar_index[i, j]] + fft_matrix[i, j] - else: - for i in range(0, nkx): - for j in range(0, nky): - for k in range(0, nkz): - fft_isok[polar_index[i, j, k]] = fft_isok[polar_index[i, j, k]] + fft_matrix[i, j, k] - - fft_isok = fft_isok / nbin[:] - return fft_isok diff --git a/src/postgkyl/tools/pressure_diagnostics.py b/src/postgkyl/tools/pressure_diagnostics.py deleted file mode 100644 index 78edc069..00000000 --- a/src/postgkyl/tools/pressure_diagnostics.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Postgkyl module for pressure tensor diagnostics. - -Diagnostics include: - Pressure parallel to the magnetic field - Pressure perpendicular to the magnetic field - Agyrotropy (either Frobenius or Swisdak measure) - Firehose instability threshold -""" - -from __future__ import annotations - -from typing import Tuple, TYPE_CHECKING -import numpy as np - -from postgkyl.tools.prim_vars import get_pij -from postgkyl.tools.mag_sq import mag_sq -from postgkyl.utils import input_parser -if TYPE_CHECKING: - from postgkyl import GData -#end - - -def _get_pb(p_in: GData | Tuple[list, np.ndarray], - b_in: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - _, p_values = input_parser(p_in) - _, b_values = input_parser(b_in) - - p_xx = p_values[..., 0, np.newaxis] - p_xy = p_values[..., 1, np.newaxis] - p_xz = p_values[..., 2, np.newaxis] - p_yy = p_values[..., 3, np.newaxis] - p_yz = p_values[..., 4, np.newaxis] - p_zz = p_values[..., 5, np.newaxis] - - b_x = b_values[..., 0, np.newaxis] - b_y = b_values[..., 1, np.newaxis] - b_z = b_values[..., 2, np.newaxis] - - return p_xx, p_xy, p_xz, p_yy, p_yz, p_zz, b_x, b_y, b_z - - -def _get_sf(species: GData | Tuple[list, np.ndarray], - field: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - p_grid, p_values = get_pij(species) - _, field_values = input_parser(field) - - b_grid = p_grid - b_values = field_values[..., 3:6] - return p_grid, p_values, b_grid, b_values - - -def get_p_par(p_in: GData | Tuple[list, np.ndarray], - b_in: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - _, p_values = input_parser(p_in) - _, b_values = input_parser(b_in) - - p_xx = p_values[..., 0, np.newaxis] - p_xy = p_values[..., 1, np.newaxis] - p_xz = p_values[..., 2, np.newaxis] - p_yy = p_values[..., 3, np.newaxis] - p_yz = p_values[..., 4, np.newaxis] - p_zz = p_values[..., 5, np.newaxis] - - b_x = b_values[..., 0, np.newaxis] - b_y = b_values[..., 1, np.newaxis] - b_z = b_values[..., 2, np.newaxis] - - grid, mag_b_sq = mag_sq(b_in) - - out = (b_x*b_x*p_xx + b_y*b_y*p_yy + b_z*b_z*p_zz - + 2.0*(b_x*b_y*p_xy + b_x*b_z*p_xz + b_y*b_z*p_yz)) / mag_b_sq - return grid, out - - -def get_gkyl_10m_p_par(species: GData | Tuple[list, np.ndarray], - field: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - p_grid, p_values = get_pij(species) - field_grid, field_values = input_parser(field) - b_values = field_values[..., 3:6] - - return get_p_par((p_grid, p_values), (field_grid, b_values)) - - -def get_p_perp(p_in: GData | Tuple[list, np.ndarray], - b_in: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - _, p_values = input_parser(p_in) - - p_xx = p_values[..., 0, np.newaxis] - p_yy = p_values[..., 3, np.newaxis] - p_zz = p_values[..., 5, np.newaxis] - - grid, p_par = get_p_par(p_in, b_in) - - out = (p_xx + p_yy + p_zz - p_par)/2.0 - return grid, out - - -def get_gkyl_10m_p_perp(species: GData | Tuple[list, np.ndarray], - field: GData | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - p_grid, p_values = get_pij(species) - field_grid, field_values = input_parser(field) - - p_grid, p_values = get_pij(species) - b_values = field_values[..., 3:6] - - return get_p_perp((p_grid, p_values), (field_grid, b_values)) - - -def get_agyro(p_in: GData | Tuple[list, np.ndarray], b_in: GData | Tuple[list, np.ndarray], - measure: str = "swisdak") -> Tuple[list, np.ndarray]: - _, p_values = input_parser(p_in) - _, b_values = input_parser(b_in) - - p_xx = p_values[..., 0, np.newaxis] - p_xy = p_values[..., 1, np.newaxis] - p_xz = p_values[..., 2, np.newaxis] - p_yy = p_values[..., 3, np.newaxis] - p_yz = p_values[..., 4, np.newaxis] - p_zz = p_values[..., 5, np.newaxis] - - b_x = b_values[..., 0, np.newaxis] - b_y = b_values[..., 1, np.newaxis] - b_z = b_values[..., 2, np.newaxis] - - grid, mag_b_sq = mag_sq(b_in) - _, p_par = get_p_par(p_in, b_in) - _, p_perp = get_p_perp(p_in, b_in) - - if measure.lower() == "swisdak": - I1 = p_xx + p_yy + p_zz - I2 = (p_xx*p_yy + p_xx*p_zz + p_yy*p_zz - - (p_xy*p_xy + p_xz*p_xz + p_yz*p_yz)) - - # Note that this definition of Q uses the tensor algebra in - # Appendix A of Swisdak 2015. - out = np.sqrt(1 - 4 * I2 / ((I1 - p_par) * (I1 + 3 * p_par))) - elif measure.lower() == "frobenius": - p_ixx = p_xx - (p_par*b_x*b_x/mag_b_sq + p_perp*(1 - b_x*b_x/mag_b_sq)) - p_ixy = p_xy - (p_par*b_x*b_y/mag_b_sq + p_perp*(0 - b_x*b_y/mag_b_sq)) - p_ixz = p_xz - (p_par*b_x*b_z/mag_b_sq + p_perp*(0 - b_x*b_z/mag_b_sq)) - p_iyy = p_yy - (p_par*b_y*b_y/mag_b_sq + p_perp*(1 - b_y*b_y/mag_b_sq)) - p_iyz = p_yz - (p_par*b_y*b_z/mag_b_sq + p_perp*(0 - b_y*b_z/mag_b_sq)) - p_izz = p_zz - (p_par*b_z*b_z/mag_b_sq + p_perp*(1 - b_z*b_z/mag_b_sq)) - out = np.sqrt(p_ixx**2 + 2*p_ixy**2 + 2*p_ixz**2 + p_iyy**2 + 2*p_iyz**2 + p_izz**2) / np.sqrt(2*p_perp**2 + 4*p_par*p_perp) - else: - raise ValueError(f"Measure specified is {measure.lower():s}; it needs to be either 'swisdak' or 'frobenius'") - # end - return grid, out - - -def get_gkyl_10m_agyro(species: GData | Tuple[list, np.ndarray], field: GData | Tuple[list, np.ndarray], - measure: str = "swisdak") -> Tuple[list, np.ndarray]: - p_grid, p_values = get_pij(species) - field_grid, field_values = input_parser(field) - b_values = field_values[..., 3:6] - - return get_agyro((p_grid, p_values), (field_grid, b_values), measure=measure) diff --git a/src/postgkyl/tools/prim_vars.py b/src/postgkyl/tools/prim_vars.py deleted file mode 100644 index d3fa19c8..00000000 --- a/src/postgkyl/tools/prim_vars.py +++ /dev/null @@ -1,396 +0,0 @@ -from __future__ import annotations - -import numpy as np -from typing import Tuple, TYPE_CHECKING - -from postgkyl.utils import input_parser -if TYPE_CHECKING: - from postgkyl import GData -# end - - -def get_density(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - out_values = in_values[..., 0, np.newaxis] - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_vx(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - out_values = in_values[..., 1, np.newaxis] / rho - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_vy(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - out_values = in_values[..., 2, np.newaxis] / rho - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_vz(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - out_values = in_values[..., 3, np.newaxis] / rho - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_vi(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - out_values = in_values[..., 1:4] / rho - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pxx(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vx = get_vx(in_mom) - out_values = in_values[..., 4, np.newaxis] - rho*vx*vx - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pxy(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vx = get_vx(in_mom) - _, vy = get_vy(in_mom) - out_values = in_values[..., 5, np.newaxis] - rho*vx*vy - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pxz(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vx = get_vx(in_mom) - _, vz = get_vz(in_mom) - out_values = in_values[..., 6, np.newaxis] - rho*vx*vz - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pyy(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vy = get_vy(in_mom) - out_values = in_values[..., 7, np.newaxis] - rho*vy*vy - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pyz(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vy = get_vy(in_mom) - _, vz = get_vz(in_mom) - out_values = in_values[..., 8, np.newaxis] - rho*vy*vz - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pzz(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vz = get_vz(in_mom) - out_values = in_values[..., 9, np.newaxis] - rho*vz*vz - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_pij(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - out_values = np.zeros(in_values[..., 4:10].shape) - - _, pxx = get_pxx(in_mom) - _, pxy = get_pxy(in_mom) - _, pxz = get_pxz(in_mom) - _, pyy = get_pyy(in_mom) - _, pyz = get_pyz(in_mom) - _, pzz = get_pzz(in_mom) - - out_values[..., 0] = np.squeeze(pxx) - out_values[..., 1] = np.squeeze(pxy) - out_values[..., 2] = np.squeeze(pxz) - out_values[..., 3] = np.squeeze(pyy) - out_values[..., 4] = np.squeeze(pyz) - out_values[..., 5] = np.squeeze(pzz) - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_p(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - num_moms: int | None = None, - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - num_comps = in_values.shape[-1] - if num_moms is None: - if num_comps == 5: - num_moms = 5 - elif num_comps == 10: - num_moms = 10 - else: - raise ValueError(f"Number of components appears to be {num_comps:d}; it needs to be specified using 'num_moms' (5 or 10)") - # end - # end - - if num_moms == 5: - _, rho = get_density(in_mom) - _, vx = get_vx(in_mom) - _, vy = get_vy(in_mom) - _, vz = get_vz(in_mom) - out_values = (gas_gamma - 1) * ( - in_values[..., 4, np.newaxis] - 0.5*rho*(vx**2 + vy**2 + vz**2) - ) - else: # num_moms == 10: - _, pxx = get_pxx(in_mom) - _, pyy = get_pyy(in_mom) - _, pzz = get_pzz(in_mom) - out_values = (pxx + pyy + pzz) / 3.0 - # end - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_ke(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - num_moms: int | None = None, - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - num_comps = in_values.shape[-1] - if num_moms is None: - if num_comps == 5: - num_moms = 5 - elif num_comps == 10: - num_moms = 10 - else: - raise ValueError(f"Number of components appears to be {num_comps:d}; (5 or 10)") - # end - # end - - if num_moms == 5: - _, pr = get_p(in_mom, gas_gamma=gas_gamma, num_moms=num_moms) - out_values = in_values[..., 4, np.newaxis] - pr / (gas_gamma - 1) - else: # num_moms == 10: - _, rho = get_density(in_mom) - _, vx = get_vx(in_mom) - _, vy = get_vy(in_mom) - _, vz = get_vz(in_mom) - out_values = 0.5*rho*(vx**2 + vy**2 + vz**2) - # end - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_temp(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - num_moms: int | None = None, - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, rho = get_density(in_mom) - _, pr = get_p(in_mom, gas_gamma=gas_gamma, num_moms=num_moms) - out_values = pr/rho - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_sound(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - num_moms: int | None = None, - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, rho = get_density(in_mom) - _, pr = get_p(in_mom, gas_gamma=gas_gamma, num_moms=num_moms) - out_values = np.sqrt(gas_gamma*pr / rho) - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mach(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - num_moms: int | None = None, - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, vx = get_vx(in_mom) - _, vy = get_vy(in_mom) - _, vz = get_vz(in_mom) - _, cs = get_sound(in_mom, gas_gamma=gas_gamma, num_moms=num_moms) - out_values = np.sqrt(vx**2 + vy**2 + vz**2) / cs - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_Bx(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - out_values = in_values[..., 5, np.newaxis] - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_By(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - out_values = in_values[..., 6, np.newaxis] - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_Bz(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - out_values = in_values[..., 7, np.newaxis] - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_Bi(in_mom: GData | Tuple[list, np.ndarray], - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - out_values = in_values[..., 5:8] - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_mag_p(in_mom: GData | Tuple[list, np.ndarray], mu_0: float = 1.0, - out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, Bx = get_mhd_Bx(in_mom) - _, By = get_mhd_By(in_mom) - _, Bz = get_mhd_Bz(in_mom) - out_values = 0.5 * (Bx**2 + By**2 + Bz**2) / mu_0 - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_p(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - mu_0: float = 1.0, out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, in_values = input_parser(in_mom) - _, rho = get_density(in_mom) - _, vx = get_vx(in_mom) - _, vy = get_vy(in_mom) - _, vz = get_vz(in_mom) - _, mag_p = get_mhd_mag_p(in_mom, mu_0=mu_0) - - out_values = (gas_gamma - 1)*(in_values[..., 4, np.newaxis] - 0.5*rho*(vx**2 + vy**2 + vz**2) - mag_p) - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_temp(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - mu_0: float = 1.0, out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, rho = get_density(in_mom) - _, pr = get_mhd_p(in_mom, gas_gamma=gas_gamma, mu_0=mu_0) - out_values = pr / rho - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_sound(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - mu_0: float = 1.0, out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, rho = get_density(in_mom) - _, pr = get_mhd_p(in_mom, gas_gamma=gas_gamma, mu_0=mu_0) - - out_values = np.sqrt(gas_gamma*pr/rho) - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values - - -def get_mhd_mach(in_mom: GData | Tuple[list, np.ndarray], gas_gamma: float = 5.0/3, - mu_0: float = 1.0, out_mom: GData | None = None) -> Tuple[list, np.ndarray]: - grid, vx = get_vx(in_mom) - _, vy = get_vy(in_mom) - _, vz = get_vz(in_mom) - _, cs = get_mhd_sound(in_mom, gas_gamma=gas_gamma, mu_0=mu_0) - out_values = np.sqrt(vx**2 + vy**2 + vz**2) / cs - - if out_mom: - out_mom.push(grid, out_values) - # end - return grid, out_values diff --git a/src/postgkyl/tools/rel_change.py b/src/postgkyl/tools/rel_change.py deleted file mode 100644 index 280f5ed8..00000000 --- a/src/postgkyl/tools/rel_change.py +++ /dev/null @@ -1,24 +0,0 @@ -import numpy as np - - -def rel_change(dataset0, dataset, comp=None): - """Function to compute the relative change in a dataset compared to another - dataset, i.e. (dataset - dataset0)/dataset0 - - Notes: - Assumes user wishes to perform this operation component-wise. - Also assumes the reference division should be performed with respect to a single - component (i.e., for energetics, divide by the total energy, - not an individual component of the energy) - """ - # Grid is the same for each of the input objects - grid = dataset.get_grid() - values = dataset.get_values() - values0 = dataset0.get_values() - out = np.zeros(values.shape) - for i in range(0, out.shape[-1]): - if comp is not None: - out[..., i] = (values[..., i] - values0[..., i]) / values0[..., int(comp)] - else: - out[..., i] = (values[..., i] - values0[..., i]) / values0[..., i] - return grid, out diff --git a/src/postgkyl/tools/rotation_matrix.py b/src/postgkyl/tools/rotation_matrix.py deleted file mode 100644 index 74cafa48..00000000 --- a/src/postgkyl/tools/rotation_matrix.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Postgkyl module including varios utility operations on fields.""" - -import numpy as np - - -def rotation_matrix(vector: np.ndarray) -> np.ndarray: - """Calculate rotation matrix. - - Args: - vector: np.ndarray - - Returns: - 3x3 rotation matrix (numpy array) - """ - rot = np.zeros((3, 3)) - norm = np.abs(vector) - k = vector / norm # direction unit vector - - # normalization - norm2 = np.sqrt(k[1]*k[1] + k[2]*k[2]) - norm3 = np.sqrt((k[1]*k[1] + k[2]*k[2])**2 + k[0]*k[0]*k[1]*k[1] + k[0]*k[0]*k[2]*k[2]) - - rot[0, :] = k - rot[1, 0] = 0 - rot[1, 1] = -k[2]/norm2 - rot[1, 2] = k[1]/norm2 - rot[2, 0] = (k[1]*k[1] + k[2]*k[2])/norm3 - rot[2, 1] = -k[0]*k[1]/norm3 - rot[2, 2] = -k[0]*k[2]/norm3 - - return rot diff --git a/src/postgkyl/tools/transform_frame.py b/src/postgkyl/tools/transform_frame.py deleted file mode 100644 index f84c1b39..00000000 --- a/src/postgkyl/tools/transform_frame.py +++ /dev/null @@ -1,95 +0,0 @@ -from __future__ import annotations - -import numpy as np -from typing import Tuple, TYPE_CHECKING - -from postgkyl.utils import input_parser -if TYPE_CHECKING: - from postgkyl import GData -# end - - -def transform_frame(in_f: GData | Tuple[list, np.ndarray], - in_u: GData | Tuple[list, np.ndarray], - c_dim: int, out_f: GData | None = None) -> Tuple[list, np.ndarray]: - """Shift a distribution function to a different frame of reference. - - Shifsts the frame of reference for specified distribution function - with a supplied bulk velocity (a direction of magnetic field will be - added in future update). - - Args: - in_f: GData or np.ndarray - Particle distribution function to be shifted. - in_u: GData or np.ndarray - Bulk velocity. - c_dim: int - Number of the configuration space dimensions. - out_f: GData - (Optional) GData to store output. - - Returns: - A tuple of grid (which is itself a tuple of nupy arrays for each - dimension) and a numpy array with values. - """ - in_f_grid, in_f_values = input_parser(in_f) - _, u = input_parser(in_u) - v_dim = len(in_f_grid) - c_dim - out_grid = np.meshgrid(*in_f_grid, indexing="ij") - - # There might be a better way to do this but hopefully such hardcoding - # is ok in this instance -- PC - if c_dim == 1: - for v_idx in range(v_dim): - nx = in_f_grid[0].shape[0] - - ext_u = np.zeros(nx) - ext_u[:-1] += u[..., v_idx] - ext_u[1:] += u[..., v_idx] - ext_u[1:-1] = ext_u[1:-1]/2 - - for i in range(nx): - out_grid[c_dim + v_idx][i, ...] += ext_u[i] - # end - # end - elif c_dim == 2: - for v_idx in range(v_dim): - nx = in_f_grid[0].shape[0] - ny = in_f_grid[0].shape[1] - - ext_u = np.zeros((nx, ny)) - ext_u[:-1, :-1] += u[..., v_idx] - ext_u[1:, 1:] += u[..., v_idx] - ext_u[1:-1, 1:-1] = ext_u[1:-1, 1:-1] / 2 - - for i in range(nx): - for j in range(ny): - out_grid[c_dim + v_idx][i, j, ...] += ext_u[i, j] - # end - # end - # end - else: - for v_idx in range(v_dim): - nx = in_f_grid[0].shape[0] - ny = in_f_grid[0].shape[1] - nz = in_f_grid[0].shape[2] - - ext_u = np.zeros((nx, ny, nz)) - ext_u[:-1, :-1, :-1] += u[..., v_idx] - ext_u[1:, 1:, 1:] += u[..., v_idx] - ext_u[1:-1, 1:-1, 1:-1] = ext_u[1:-1, 1:-1, 1:-1]/2 - - for i in range(nx): - for j in range(ny): - for k in range(nz): - out_grid[c_dim + v_idx][i, j, k, ...] += ext_u[i, j, k] - # end - # end - # end - # end - # end - - if out_f: - out_f.push(out_grid, in_f_values) - # end - return out_grid, in_f_values diff --git a/src/postgkyl/utils/__init__.py b/src/postgkyl/utils/__init__.py deleted file mode 100644 index 8461bc26..00000000 --- a/src/postgkyl/utils/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .input_parser import input_parser -from .load_style import load_style -from .verb_print import verb_print -from .set_frame import set_frame diff --git a/src/postgkyl/utils/gk_quantities/fetch_funcs.py b/src/postgkyl/utils/gk_quantities/fetch_funcs.py deleted file mode 100644 index 3736f89f..00000000 --- a/src/postgkyl/utils/gk_quantities/fetch_funcs.py +++ /dev/null @@ -1,947 +0,0 @@ -""" -Functions for for fetching (loading and computing) quantities in the -gk_quantities registry. - -Each fetch function takes a list of loaded GData objects (matching the -corresponding 'files' entry in the registry) and returns (grid, values) for -the derived quantity. - -Naming keys for some fetch functions below: - s#: source # - c#: component # - add: plus - sub: minus - mul: times - div: divided by - pow#: raised to the power of # - -""" -import numpy as np -import operator - -from postgkyl.data import GData -from postgkyl.data.dg import get_num_basis -from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops -import postgkyl.utils.gkeyll_const as gkc - -def _get_ctx_val(gdata : GData, key : str, **kwargs): - """ - Read a value(s) for 'key', the '--extra' value overrides the GData's context. - """ - if key in kwargs: - val = kwargs[key] - - if not isinstance(val, (list, tuple)): - # A single value applies to every species. - return val - - species_idx = kwargs.get("species_idx", None) - if species_idx is None: - raise KeyError(f"fetch function: '--extra {key}=' was given {len(val)} values but this " - f"quantity is not computed per species, so there is no way to tell which " - f"one to use. Pass a single value instead.") - - if species_idx >= len(val): - species = kwargs.get("species", None) - raise ValueError(f"fetch function: '--extra {key}=' was given only {len(val)} values but " - f"species #{species_idx}{f' ({species})' if species else ''} was requested. " - f"Give one value per species, in the order of '--species'.") - - return val[species_idx] - - if gdata.ctx.get(key, None) is not None: - return gdata.ctx[key] - - raise KeyError(f"fetch function: context key '{key}' not found in GData. Pass it as " - f"'--extra {key}=', or as one value per species with " - f"'--extra {key}=,,...'.") - -def _get_num_basis_from_gdata(gdata) -> int: - from postgkyl.data.dg import get_num_basis - ndim = gdata.get_num_dims() - poly_order = int(gdata.ctx["poly_order"]) - basis_type = gdata.ctx["basis_type"] - return get_num_basis(ndim, poly_order, basis_type) - -def _empty_gdata_from_gdata(gdata) -> GData: - """Allocate a zero-valued GData with the same grid/ctx as gdata.""" - out = GData(ctx=gdata.ctx) - out.push(gdata.get_grid(), np.zeros_like(gdata.get_values())) - return out - -def _powsqrt_dg(gdata, exponent: float) -> GData: - """ - pow(sqrt(f), exponent) of a single-component DG field. negative values are set to 1e-40. - """ - - out = _empty_gdata_from_gdata(gdata) - - dgops = GkeyllDGops() - dgops.powsqrt(out, gdata, exponent) - - return out - -def _make_fetch_comp(icomp: int): - """Return a fetch function that extracts the comp-th physical component.""" - def fetch(gdatas, **kw): - g = gdatas[0].get_grid() - nb = _get_num_basis_from_gdata(gdatas[0]) - comp = [icomp,icomp] if icomp is not None else [0,int(gdatas[0].get_num_comps()/nb)] - v = gdatas[0].get_values()[..., comp[0]*nb:(comp[1]+1)*nb].copy() - out = GData(ctx=gdatas[0].ctx) - out.push(g, v) - return out - # end - fetch.__name__ = f"fetch_comp{icomp}" if icomp is not None else f"fetch_compAll" - return fetch - -def _make_fetch_sick_addsub_sjcl(si: int, ck: int, sj: int, cl: int, op): - """ - Return a fetch function that does: - (k-th component of the i-th source) op (l-th component of the j-th source) - """ - def fetch(gdatas, **kwargs): - gd_l = gdatas[si] - gd_r = gdatas[sj] - - nb_l = _get_num_basis_from_gdata(gd_l) - nb_r = _get_num_basis_from_gdata(gd_r) - if not nb_l == nb_r: - raise ValueError(f"Datasets have different basis") - - vals_l = gd_l.get_values()[..., ck*nb_l:(ck+1)*nb_l] - vals_r = gd_r.get_values()[..., cl*nb_r:(cl+1)*nb_r] - - out = GData(ctx=gd_l.ctx) - out.push(gd_l.get_grid(), op(vals_l,vals_r)) - - return out - # end - fetch.__name__ = f"fetch_s{si}c{ck}_{op.__name__}_s{sj}c{cl}" - return fetch - -def _make_fetch_sick_mul_sjcl(si: int, ck: int, sj: int, cl: int): - """ - Return a fetch function that multiplies the k-th component of the i-th - source/dataset by the l-th component of the j-th source. - """ - def fetch(gdatas, **kwargs): - gd_l = gdatas[si] - gd_r = gdatas[sj] - - nb_l = _get_num_basis_from_gdata(gd_l) - nb_r = _get_num_basis_from_gdata(gd_r) - if not nb_l == nb_r: - raise ValueError(f"Datasets have different basis") - - vals_l = gd_l.get_values() - out_shape = list(vals_l.shape) - out_shape[-1] = nb_l - - out = GData(ctx=gd_l.ctx) - out.push(gd_l.get_grid(), np.zeros(out_shape, dtype=vals_l.dtype)) - - dgops = GkeyllDGops() - dgops.multiply(0, out, ck, gd_l, cl, gd_r) - - return out - # end - fetch.__name__ = f"fetch_s{si}c{ck}_mul_s{sj}c{cl}" - return fetch - -def _make_fetch_sick_div_sjcl(si: int, ck: int, sj: int, cl: int): - """ - Return a fetch function that divides the k-th component of the i-th - source/dataset by the l-th component of the j-th source. - """ - def fetch(gdatas, **kwargs): - gd_l = gdatas[si] - gd_r = gdatas[sj] - - nb_l = _get_num_basis_from_gdata(gd_l) - nb_r = _get_num_basis_from_gdata(gd_r) - if not nb_l == nb_r: - raise ValueError(f"Datasets have different basis") - - vals_l = gd_l.get_values() - out_shape = list(vals_l.shape) - out_shape[-1] = nb_l - - out = GData(ctx=gd_l.ctx) - out.push(gd_l.get_grid(), np.zeros(out_shape, dtype=vals_l.dtype)) - - dgops = GkeyllDGops() - dgops.invert(0, out, cl, gd_r) - dgops.multiply(0, out, ck, gd_l, 0, out) - - return out - # end - fetch.__name__ = f"fetch_s{si}c{ck}_div_s{sj}c{cl}" - return fetch - -def _b_cross_grad_div_B_component(scalar, jacobtot_inv, b_i, comp): - """ - The comp-th component of the cross product b x grad(f) - (b x grad f)_k / B = epsilon_{ijk} * b_i * d(f)/dx^j / (J B) - where epsilon_{ijk} is the Levi-Civitta tensor, f is a scalar field - and b_i are the covariant components of a vector field. - - Note: the 1/Jacobian factor of the curvilinear cross product is NOT - included here and must be applied by the caller. - - Inputs: - scalar: scalar field f to be differentiated. - jacobtot_inv: inverse of the Jacobian of the total coordinate transformation. - b_i: covariant components of the vector field b. - comp: component k of the cross product to compute (0-index, < 3). - """ - cdim = scalar.get_num_dims() - - # Components of the quantities in the cross product AxB. - diff_dir_pos = bi_c_pos = 0 - diff_dir_neg = bi_c_neg = 0 - calc_term = [True,True] # Whether to compute pos and neg term in component of AxB. - if comp == 0: - diff_dir_neg = bi_c_pos = 1 - diff_dir_pos = bi_c_neg = cdim-1 - if cdim < 3: - calc_term = [True,False] - # end - elif comp == 1: - bi_c_pos = 2 - bi_c_neg = 0 - diff_dir_neg = cdim-1 - diff_dir_pos = 0 - if cdim == 1: - calc_term = [False,True] - # end - elif comp == 2: - diff_dir_neg = bi_c_pos = 0 - diff_dir_pos = bi_c_neg = 1 - if cdim == 1: - calc_term = [False,False] - elif cdim == 2: - calc_term = [False,True] - # end - else: - raise KeyError("_b_cross_grad_component: component must be >= 0 and < 3.") - - buff = _empty_gdata_from_gdata(scalar) # Positive term in AxB. - out = _empty_gdata_from_gdata(scalar) # Negative term in AxB. - - dgops = GkeyllDGops() - lower, upper = scalar.get_bounds() - cells = scalar.get_num_cells() - if calc_term[0]: - # Compute derivatives of the scalar field. - dx = (upper[diff_dir_pos] - lower[diff_dir_pos])/cells[diff_dir_pos] - dgops.differentiate(diff_dir_pos, 1, dx, 0, buff, 0, scalar) - # Multiply by b_i. - dgops.multiply(0, buff, bi_c_pos, b_i, 0, buff) - - if calc_term[1]: - # Compute derivatives of the scalar field. - dx = (upper[diff_dir_neg] - lower[diff_dir_neg])/cells[diff_dir_neg] - dgops.differentiate(diff_dir_neg, 1, -dx, 0, out , 0, scalar) - # Multiply by b_i. - dgops.multiply(0, out , bi_c_neg, b_i, 0, out ) - - # Add the two terms to form the comp-th component of b x grad(f). - out.set_values(buff.get_values() + out.get_values()) - - # Divide by the Jacobian factor of the curvilinear cross product. - dgops.multiply(0, out, 0, out, 0, jacobtot_inv) - - return out - -# Functions to extract a components. -fetch_s0cAll = _make_fetch_comp(None) -fetch_s0c0 = _make_fetch_comp(0) -fetch_s0c1 = _make_fetch_comp(1) -fetch_s0c2 = _make_fetch_comp(2) -fetch_s0c3 = _make_fetch_comp(3) - -# Functions to add two components. -fetch_s0c0_add_s1c0 = _make_fetch_sick_addsub_sjcl(0,0,1,0,operator.add) -fetch_s0c2_add_s0c3 = _make_fetch_sick_addsub_sjcl(0,2,0,3,operator.add) - -# Functions to subtract two components. -fetch_s0c0_sub_s1c0 = _make_fetch_sick_addsub_sjcl(0,0,1,0,operator.sub) - -# Functions to multiply two components. -fetch_s0c0_mul_s1c0 = _make_fetch_sick_mul_sjcl(0,0,1,0) -fetch_s0c0_mul_s0c1 = _make_fetch_sick_mul_sjcl(0,0,0,1) - -# Functions to divide two components. -fetch_s1c0_div_s0c0 = _make_fetch_sick_div_sjcl(1,0,0,0) - -# ------------------------------------------ -# --- Plasma moments (species-dependent) --- -# ------------------------------------------ - -def fetch_M1_from_H(gdatas, **kwargs): - """ - M1 from the Hamiltonian moments (Hmom). - """ - hmom = gdatas[0] - mass = _get_ctx_val(hmom, "mass", **kwargs) - nb = _get_num_basis_from_gdata(hmom) - vals = hmom.get_values() - - m1 = GData(ctx=hmom.ctx) - m1.push(hmom.get_grid(), np.zeros_like(vals[..., :nb])) - - dgops = GkeyllDGops() - dgops.multiply(0, m1, 0, hmom, 1, hmom) - - m1.set_values(m1.get_values() / mass) - return m1 - -def fetch_Tpar_from_BiMax(gdatas, **kwargs): - """ - Tpar from BiMaxwellian moments. - """ - Tpar = fetch_s0c2(gdatas) - - bimax = gdatas[0] - mass = _get_ctx_val(bimax, "mass", **kwargs) - Tpar.set_values(mass * Tpar.get_values()) - return Tpar - -def fetch_Tpar_from_M0_M1_M2par(gdatas, **kwargs): - """ - upar*M1 + M0*Tpar/m = M2par. - Tpar = m * (M2par - upar*M1) / M0. - """ - m0, m1, m2par = gdatas - dgops = GkeyllDGops() - - m0_inv = _empty_gdata_from_gdata(m0) - upar = _empty_gdata_from_gdata(m0) - Tpar = _empty_gdata_from_gdata(m0) - - dgops.invert(0, m0_inv, 0, m0) - dgops.multiply(0, upar, 0, m1, 0, m0_inv) - dgops.multiply(0, upar, 0, upar, 0, m1) - - m2par_val = m2par.get_values() - um1_val = upar.get_values() - - mass = _get_ctx_val(m0, "mass", **kwargs) - Tpar.set_values(mass * (m2par_val - um1_val)) - dgops.multiply(0, Tpar, 0, Tpar, 0, m0_inv) - return Tpar - -def fetch_Tperp_from_BiMax(gdatas, **kwargs): - """ - Tperp from BiMaxwellian moments. - """ - Tperp = fetch_s0c3(gdatas) - - bimax = gdatas[0] - mass = _get_ctx_val(bimax, "mass", **kwargs) - Tperp.set_values(mass * Tperp.get_values()) - return Tperp - -def fetch_Tperp_from_M0_M2perp(gdatas, **kwargs): - """ - Tperp = 0.5 * mass * (M2perp / M0). - """ - Tperp = fetch_s1c0_div_s0c0(gdatas) - - m0 = gdatas[0] - mass = _get_ctx_val(m0, "mass", **kwargs) - Tperp.set_values(0.5 * mass * Tperp.get_values()) - return Tperp - -def fetch_temp_from_Max(gdatas, **kwargs): - """ - temp from Maxwellian moments. - """ - temp = fetch_s0c2(gdatas) - - maxmom = gdatas[0] - mass = _get_ctx_val(maxmom, "mass", **kwargs) - temp.set_values(mass * temp.get_values()) - return temp - -def fetch_temp_from_Tpar_Tperp(gdatas, **kwargs): - """ - temp = (Tpar + 2*Tperp) / 3. - """ - Tpar, Tperp = gdatas - - temp = _empty_gdata_from_gdata(Tpar) - - Tpar_val = Tpar.get_values() - Tperp_val = Tperp.get_values() - - temp.set_values((Tpar_val + 2.0*Tperp_val)/3.0) - return temp - -# --------------------------------------------------- -# --- Combined plasma moments (species-dependent) --- -# --------------------------------------------------- - -def fetch_press_from_Max(gdatas, **kwargs): - """ - Pressure from Maxwellian moments. - press = den * temp. - """ - maxmom = gdatas[0] - nb = _get_num_basis_from_gdata(maxmom) - vals = maxmom.get_values()[..., :nb] - - press = GData(ctx=maxmom.ctx) - press.push(maxmom.get_grid(), np.zeros_like(vals)) - - dgops = GkeyllDGops() - dgops.multiply(0, press, 0, maxmom, 2, maxmom) - - mass = _get_ctx_val(maxmom, "mass", **kwargs) - press.set_values(mass * press.get_values()) - return press - -def fetch_press_from_BiMax(gdatas, **kwargs): - """ - Pressure from BiMaxwellian moments. - press = den * (Tpar + 2*Tperp) / 3. - """ - bimax = gdatas[0] - nb = _get_num_basis_from_gdata(bimax) - vals = bimax.get_values() - - mass = _get_ctx_val(bimax, "mass", **kwargs) - Tpar_vals = vals[..., 2*nb:3*nb] - Tperp_vals = vals[..., 3*nb:4*nb] - temp_vals = mass*(Tpar_vals + 2.0 * Tperp_vals)/3.0 - - press = GData(ctx=bimax.ctx) - press.push(bimax.get_grid(), temp_vals.copy()) - - dgops = GkeyllDGops() - dgops.multiply(0, press, 0, bimax, 0, press) - - return press - -def fetch_press_p(gdatas, **kwargs): - """ - Perpendicular/parallel pressure in J/m^3. - p_p = n * T_p. - """ - m0 = gdatas[0] - Tp = gdatas[1] - - dgops = GkeyllDGops() - press_p = _empty_gdata_from_gdata(m0) - dgops.multiply(0, press_p, 0, m0, 0, Tp) - - return press_p - -def _make_fetch_q(name: str): - """ - Return a fetch function for the lab-frame parallel flux of the parallel - (name='par') or perpendicular (name='perp') kinetic energy: - q_par = (m/2)*M3par = (m/2) int(vpar^3 f) dv, - q_perp = (m/2)*M3perp = (m/2) int(vpar*vperp^2 f) dv, - so that q_par + q_perp is the parallel flux of the total kinetic energy. - Both are in W/m^2 (kg/s^3). gdatas has: - 1. M3par (name='par') or M3perp (name='perp'). - """ - def fetch(gdatas, **kwargs): - m3 = gdatas[0] - mass = _get_ctx_val(m3, "mass", **kwargs) - - out = _empty_gdata_from_gdata(m3) - out.set_values(0.5*mass*m3.get_values()) - return out - # end - fetch.__name__ = f"fetch_q{name}" - return fetch - -fetch_qpar = _make_fetch_q("par") -fetch_qperp = _make_fetch_q("perp") - -def _make_fetch_q_fluid(name: str): - """ - Return a fetch function for the parallel heat flux in the fluid (drift) - frame, i.e. the energy carried by the random part of the parallel motion, - u = M1/M0 being the parallel drift speed: - q_par = (m/2) int (vpar-u)^3 f dv - = (m/2) [M3par - 3*u*M2par + 3*u^2*M1 - u^3*M0] - = (m/2) [M3par - 3*u*M2par + 2*u^2*M1], - q_perp = (m/2) int (vpar-u)*vperp^2 f dv - = (m/2) [M3perp - u*M2perp]. - gdatas has (in this order): - 1. M0: zeroth moment (density). - 2. M1: first moment. - 3. M2par (name='par') or M2perp (name='perp'). - 4. M3par (name='par') or M3perp (name='perp'). - """ - is_par = name == "par" - - def fetch(gdatas, **kwargs): - m0, m1, m2, m3 = gdatas - mass = _get_ctx_val(m0, "mass", **kwargs) - - dgops = GkeyllDGops() - - m0_inv = _empty_gdata_from_gdata(m0) - dgops.invert(0, m0_inv, 0, m0) - - upar = _empty_gdata_from_gdata(m0) - dgops.multiply(0, upar, 0, m1, 0, m0_inv) - - # u*M2par or u*M2perp. - u_m2 = _empty_gdata_from_gdata(m0) - dgops.multiply(0, u_m2, 0, upar, 0, m2) - - if is_par: - # u^2*M1, which equals u^3*M0. - u_sq = _empty_gdata_from_gdata(m0) - dgops.multiply(0, u_sq, 0, upar, 0, upar) - - u_sq_m1 = _empty_gdata_from_gdata(m0) - dgops.multiply(0, u_sq_m1, 0, u_sq, 0, m1) - - vals = m3.get_values() - 3.0*u_m2.get_values() + 2.0*u_sq_m1.get_values() - else: - vals = m3.get_values() - u_m2.get_values() - - out = _empty_gdata_from_gdata(m0) - out.set_values(0.5*mass*vals) - return out - - fetch.__name__ = f"fetch_q{name}_fluid" - return fetch - -fetch_qpar_fluid = _make_fetch_q_fluid("par") -fetch_qperp_fluid = _make_fetch_q_fluid("perp") - -def fetch_vt(gdatas, **kwargs): - """ - Thermal speed vt = sqrt(T/m) (m/s), where T is the temperature of the - requested species and m its mass. gdatas has: - 1. temp: temperature (in Joules). - """ - temp = gdatas[0] - mass = _get_ctx_val(temp, "mass", **kwargs) - - temp_over_m = _empty_gdata_from_gdata(temp) - temp_over_m.set_values(temp.get_values()/mass) - - return _powsqrt_dg(temp_over_m, 1.0) - -def fetch_larmor_radius(gdatas, **kwargs): - """ - Species Larmor (gyro-)radius: rho = sqrt(m*T)/(|q|*B). gdatas has: - 1. temp: temperature (in Joules). - 2. Bmag: magnetic field magnitude (bmag). - """ - temp, bmag = gdatas - mass = _get_ctx_val(temp, "mass", **kwargs) - charge = abs(_get_ctx_val(temp, "charge", **kwargs)) - - mT = _empty_gdata_from_gdata(temp) - mT.set_values(temp.get_values() * mass) - sqrt_mT = _powsqrt_dg(mT, 1.0) - - qB = _empty_gdata_from_gdata(bmag) - qB.set_values(bmag.get_values() * charge) - - dgops = GkeyllDGops() - - qB_inv = _empty_gdata_from_gdata(bmag) - dgops.invert(0, qB_inv, 0, qB) - - out = _empty_gdata_from_gdata(bmag) - dgops.multiply(0, out, 0, sqrt_mT, 0, qB_inv) - - return out - -def fetch_debye_length(gdatas, **kwargs): - """ - Species-wise Debye length: lambda_D = sqrt(eps0*T/(n*q^2)). gdatas has: - 1. temp: temperature (in Joules). - 2. M0: zeroth moment (density). - """ - temp, m0 = gdatas - charge = _get_ctx_val(temp, "charge", **kwargs) - eps0 = gkc.GKYL_EPSILON0 - - eps0T = _empty_gdata_from_gdata(temp) - eps0T.set_values(temp.get_values() * eps0) - - nq2 = _empty_gdata_from_gdata(m0) - nq2.set_values(m0.get_values() * charge**2) - - dgops = GkeyllDGops() - - nq2_inv = _empty_gdata_from_gdata(m0) - dgops.invert(0, nq2_inv, 0, nq2) - - sq = _empty_gdata_from_gdata(temp) - dgops.multiply(0, sq, 0, eps0T, 0, nq2_inv) - - return _powsqrt_dg(sq, 1.0) - -def _split_elc_ions(gdatas, quantity: str, **kwargs): - """ - Split the per-species sources of a multi-species quantity into the electron - entry and the ion entries, by the sign of each species' charge.. - """ - species_names = kwargs.get("species", []) - if len(species_names) != len(gdatas): - species_names = [f"#{i}" for i in range(len(gdatas))] - - elcs, ions = [], [] - for species_idx, (name, srcs) in enumerate(zip(species_names, gdatas)): - # Resolve each species' attributes against its own slot in a '--extra' array. - species_kwargs = dict(kwargs, species_idx=species_idx, species=name) - entry = { - "name": name, - "srcs": srcs, - "mass": _get_ctx_val(srcs[0], "mass", **species_kwargs), - "charge": _get_ctx_val(srcs[0], "charge", **species_kwargs), - } - (elcs if entry["charge"] < 0.0 else ions).append(entry) - - if len(elcs) != 1: - raise ValueError(f"{quantity}: expected exactly one negatively charged (electron) species " - f"but found {len(elcs)} in {list(species_names)}.") - - if not ions: - raise ValueError(f"{quantity}: found no positively charged (ion) species in {list(species_names)}.") - - return elcs[0], ions - -def _weighted_sum(entries, weights, comp: int): - """ - Sum the comp-th source of each species, each scaled by a scalar weight. - """ - out = _empty_gdata_from_gdata(entries[0]["srcs"][comp]) - total = sum(w*e["srcs"][comp].get_values() for e, w in zip(entries, weights)) - out.set_values(total) - return out - -def _fetch_c_s_ion_acoustic(gdatas, **kwargs): - """ - Ion-acoustic sound speed (wave perspective), for the Bohm criterion and - sheath/presheath matching: - c_s = sqrt( T_e * sum_j(n_j*Z_j^2/m_j) / sum_j(n_j*Z_j) ) - summing over the ion species j, with Z_j = q_j/e the ion charge state. - """ - elc, ions = _split_elc_ions(gdatas, "fetch_c_s(kind=ion_acoustic)", **kwargs) - - e = gkc.GKYL_ELEMENTARY_CHARGE - charge_states = [ion["charge"]/e for ion in ions] - - # sum_j n_j*Z_j^2/m_j and sum_j n_j*Z_j, both linear in the densities (M0). - numer = _weighted_sum(ions, [z**2/ion["mass"] for z, ion in zip(charge_states, ions)], 0) - denom = _weighted_sum(ions, charge_states, 0) - - dgops = GkeyllDGops() - - denom_inv = _empty_gdata_from_gdata(denom) - dgops.invert(0, denom_inv, 0, denom) - - # T_e * numer/denom. - c_s_sq = _empty_gdata_from_gdata(numer) - dgops.multiply(0, c_s_sq, 0, numer, 0, denom_inv) - dgops.multiply(0, c_s_sq, 0, c_s_sq, 0, elc["srcs"][1]) - - return _powsqrt_dg(c_s_sq, 1.0) - -def _fetch_c_s_thermo(gdatas, **kwargs): - """ - Thermodynamic sound speed (bulk fluid perspective), for Mach numbers and - acoustic propagation in the core/SOL: - c_s = sqrt( (gamma_e*n_e*T_e + sum_j(gamma_j*n_j*T_j)) / sum_j(n_j*m_j) ) - summing over the ion species j. - Default: gamma_e=1, gamma_i=3, but these can be set via '--extra'. - """ - elc, ions = _split_elc_ions(gdatas, "fetch_c_s(kind=thermo)", **kwargs) - - gamma_e = float(kwargs.get("gamma_e", 1.0)) - gamma_i = float(kwargs.get("gamma_i", 3.0)) - - dgops = GkeyllDGops() - - # gamma_e*n_e*T_e + sum_j gamma_j*n_j*T_j. Each n*T is a weak product. - numer = _empty_gdata_from_gdata(elc["srcs"][0]) - dgops.multiply(0, numer, 0, elc["srcs"][0], 0, elc["srcs"][1]) - numer.set_values(gamma_e*numer.get_values()) - - press_j = _empty_gdata_from_gdata(elc["srcs"][0]) - for ion in ions: - dgops.multiply(0, press_j, 0, ion["srcs"][0], 0, ion["srcs"][1]) - numer.set_values(numer.get_values() + gamma_i*press_j.get_values()) - - # sum_j n_j*m_j, the ion mass density; linear in the densities. - denom = _weighted_sum(ions, [ion["mass"] for ion in ions], 0) - - denom_inv = _empty_gdata_from_gdata(denom) - dgops.invert(0, denom_inv, 0, denom) - - c_s_sq = _empty_gdata_from_gdata(numer) - dgops.multiply(0, c_s_sq, 0, numer, 0, denom_inv) - - return _powsqrt_dg(c_s_sq, 1.0) - -def fetch_c_s(gdatas, **kwargs): - """ - Sound speed (m/s), combining the electrons and every ion species. gdatas has - one [M0, temp] pair per species, in the order they were requested, e.g. - pgkyl gk-load-quantity -q c_s -s elc,ion1,ion2 ... - Electrons and ions are told apart by the sign of each species' charge - attribute, so the species may be named anything. - - Two definitions are available through '--extra kind=': - ion_acoustic (default): the wave/Bohm-criterion sound speed, - c_s = sqrt(T_e*sum_j(n_j*Z_j^2/m_j)/sum_j(n_j*Z_j)). - thermo: the bulk-fluid sound speed, - c_s = sqrt((gamma_e*n_e*T_e + sum_j(gamma_j*n_j*T_j))/sum_j(n_j*m_j)), - with gamma_e and gamma_i settable via '--extra' (default 1 and 3). - """ - c_s_kinds = { - "ion_acoustic": _fetch_c_s_ion_acoustic, - "thermo": _fetch_c_s_thermo, - } - kind = str(kwargs.get("kind", "thermo")) - if kind not in c_s_kinds: - raise ValueError(f"fetch_c_s: unknown kind '{kind}'. Select one with '--extra kind=' " - f"from: {', '.join(sorted(c_s_kinds))}.") - # end - return c_s_kinds[kind](gdatas, **kwargs) - -def fetch_beta_from_bmag_press(gdatas, **kwargs): - """ - beta = 2*mu_0*press/bmag^2 - """ - bmag, press = gdatas - - dgops = GkeyllDGops() - - bmag_sq = _empty_gdata_from_gdata(bmag) - out = _empty_gdata_from_gdata(bmag) - - dgops.multiply(0, bmag_sq, 0, bmag, 0, bmag) - - dgops.invert(0, out, 0, bmag_sq) - dgops.multiply(0, out, 0, press, 0, out) - - out_val = out.get_values() - - mu0 = gkc.GKYL_MU0 - out.set_values(2.0*mu0*out_val) - return out - -# ------------------------ -# --- Drift velocities --- -# ------------------------ - -def fetch_ExB_vel(gdatas, **kwargs): - """ - A component of the ExB drift velocity - v_{E,k} = epsilon_{ijk}/(J B) * b_i * d(phi)/dx^j - where epsilon_{ijk} is the Levi-Civitta tensor - and gdatas has (in this order): - 1/(J*B): jacobtot_inv. - b_i: covariant components of the magnetic field unit vector. - phi: electrostatic potential. - - The k-th component is selected by the 'dir' optional argument. - """ - if "dir" not in kwargs: - raise KeyError("fetch_ExB_vel: select the j-th component with '--extra dir=j' (0-index).") - - jacobtot_inv = gdatas[0] - bmag = gdatas[1] - b_i = gdatas[2] - phi = gdatas[3] - - # k-th component of b x grad(phi)/B. - out = _b_cross_grad_div_B_component(phi, jacobtot_inv, b_i, kwargs["dir"]) - - return out - -def fetch_gradB_vel(gdatas, **kwargs): - """ - A component of the grad-B drift velocity - v_gradB,k = Tperp/(q B) * epsilon_{ijk} * b_i * d(B)/dx^j / (J B) - where epsilon_{ijk} is the Levi-Civitta tensor, q the species charge, - and gdatas has (in this order): - 1/(J*B): inv. total Jacobian (jacobtot_inv). - B: magnetic field magnitude (bmag). - b_i: covariant components of the magnetic field unit vector. - Tperp: perpendicular temperature (in Joules). - - The k-th component is selected by the 'dir' optional argument. - """ - if "dir" not in kwargs: - raise KeyError("fetch_gradB_vel: select the j-th component with '--extra dir=j' (0-index).") - - jacobtot_inv = gdatas[0] - bmag = gdatas[1] - b_i = gdatas[2] - Tperp = gdatas[3] - - # k-th component of b x grad(B)/B. - out = _b_cross_grad_div_B_component(bmag, jacobtot_inv, b_i, kwargs["dir"]) - - dgops = GkeyllDGops() - # Multiply by Tperp. - dgops.multiply(0, out, 0, Tperp, 0, out) - - # Divide by B. - denom_inv = _empty_gdata_from_gdata(bmag) - dgops.invert(0, denom_inv, 0, bmag) - dgops.multiply(0, out, 0, out, 0, denom_inv) - - # Divide by the species charge. - charge = _get_ctx_val(Tperp, "charge", **kwargs) - out.set_values(out.get_values()/charge) - - return out - -def fetch_diamag_vel(gdatas, **kwargs): - """ - A component of the diamagnetic drift velocity - v_diamag,k = 1 / (q n) epsilon_{ijk} b_i * d(pperp)/dx^j / (J B) - where epsilon_{ijk} is the Levi-Civitta tensor, q the species charge, - and gdatas has (in this order): - 1/(J*B): inv. total Jacobian (jacobtot_inv). - B: magnetic field magnitude (bmag). - b_i: covariant components of the magnetic field unit vector. - m0: zeroth moment (density). - p_perp: perpendicular pressure (in Joules/m^3). - The k-th component is selected by the 'dir' optional argument. - """ - if "dir" not in kwargs: - raise KeyError("fetch_diamag_vel: select the j-th component with '--extra dir=j' (0-index).") - - jacobtot_inv = gdatas[0] - bmag = gdatas[1] - b_i = gdatas[2] - m0 = gdatas[3] - pressperp = gdatas[4] - - # k-th component of b x grad(p) / B. - out = _b_cross_grad_div_B_component(pressperp, jacobtot_inv, b_i, kwargs["dir"]) - - dgops = GkeyllDGops() - # Divide by n - denom_inv = _empty_gdata_from_gdata(bmag) - dgops.invert(0, denom_inv, 0, m0) - dgops.multiply(0, out, 0, out, 0, denom_inv) - - # Divide by the species charge. - charge = _get_ctx_val(pressperp, "charge", **kwargs) - out.set_values(out.get_values()/charge) - - return out - -def load_distf(gdatas, **kwargs) -> GData: - """ - Loader for the registry 'distf' quantity. Wraps load_gk_distf with defaults - tailored to registry use: never interpolate (interp=0) and convert velocity - coordinates (c2p_vel) on by default. - - Defaults can be overridden via --extra, e.g.: - -e suffix=source use -_source_.gkyl as input - -e c2p_vel=0 disable velocity-space mapping - -e mc2nu=1 apply non-uniform -> field-aligned position mapping - -e mapc2p=1 apply position-space -> Cartesian/cylindrical mapping - -e block=2 load only the 2nd block of a multi-block file - """ - from postgkyl.commands.gk_distf import load_gk_distf - from postgkyl.utils.gk_utils import dict_get_bool - - prefix = kwargs.get("path", "").rstrip("/") + "/" + kwargs.get("name", "") - extra = kwargs.get("extra", {}) - - return load_gk_distf( - name=prefix, species=kwargs.get("species", ""), frame=int(kwargs.get("frame", 0)), - suffix=str(extra.get("suffix", "")), - use_c2p_vel=dict_get_bool(extra, "c2p_vel", True), - use_mc2nu=dict_get_bool(extra, "mc2nu", False), - use_mapc2p=dict_get_bool(extra, "mapc2p", False), - block_idx=extra.get("block", None), - interp=0, # registry distf always works with non-interpolated DG data - ) - -def _make_fetch_q_norm(name: str): - """ - Return a fetch function for a heat flux normalized by the free-streaming - estimate n*T*c_s: - q_norm = q / (n*T*c_s). - gdatas has (in this order): - 1. q: the heat flux to normalize (in W/m^2). - 2. M0: zeroth moment (density). - 3. temp: temperature (in Joules). - 4. c_s: sound speed (in m/s). - """ - def fetch(gdatas, **kwargs): - q, m0, temp, c_s = gdatas - - dgops = GkeyllDGops() - - # n*T*c_s. - denom = _empty_gdata_from_gdata(m0) - dgops.multiply(0, denom, 0, m0, 0, temp) - dgops.multiply(0, denom, 0, denom, 0, c_s) - - denom_inv = _empty_gdata_from_gdata(m0) - dgops.invert(0, denom_inv, 0, denom) - - out = _empty_gdata_from_gdata(m0) - dgops.multiply(0, out, 0, q, 0, denom_inv) - return out - - fetch.__name__ = f"fetch_q{name}_norm" - return fetch - -fetch_qpar_norm = _make_fetch_q_norm("par") -fetch_qperp_norm = _make_fetch_q_norm("perp") - - -def fetch_rho_over_lambda(gdatas, **kwargs): - """ - Ratio of the species Larmor radius to its Debye length: rho/lambda_D. - gdatas has: - 1. rho: Larmor radius (m). - 2. lambda_D: Debye length (m). - """ - rho, lambda_d = gdatas - - dgops = GkeyllDGops() - - lambda_d_inv = _empty_gdata_from_gdata(lambda_d) - dgops.invert(0, lambda_d_inv, 0, lambda_d) - - out = _empty_gdata_from_gdata(rho) - dgops.multiply(0, out, 0, rho, 0, lambda_d_inv) - - return out - -def fetch_phi_norm(gdatas, **kwargs): - """ - Normalized electrostatic potential. - phi_norm = e*phi/T_e. Gdatas has: - 1. phi: electrostatic potential (phi). - 2. temp: temperature (temp). - """ - phi, temp = gdatas - e = gkc.GKYL_ELEMENTARY_CHARGE - - dgops = GkeyllDGops() - - temp_inv = _empty_gdata_from_gdata(temp) - dgops.invert(0, temp_inv, 0, temp) - - out = _empty_gdata_from_gdata(phi) - dgops.multiply(0, out, 0, phi, 0, temp_inv) - - out.set_values(out.get_values() * e) - - return out \ No newline at end of file diff --git a/src/postgkyl/utils/gk_quantities/gkquantity.py b/src/postgkyl/utils/gk_quantities/gkquantity.py deleted file mode 100644 index 904341c1..00000000 --- a/src/postgkyl/utils/gk_quantities/gkquantity.py +++ /dev/null @@ -1,306 +0,0 @@ -import glob -import os - -from postgkyl.data import GData - -class GkQuantity: - """ - Class for a gyrokinetic quantity. - - Attributes: - name: Name of the quantity. - source: List of file combinations to try. - fetch_func: Corresponding fetch function for each file combo. - label: LaTeX format label for matplotlib (use %s for species name or direction). - is_time_dep: If the quantity is time-dependent (i.e. written in frames). - is_species_dep: If the quantity is species-dependent. - is_vector: If the quantity is a vector (i.e. has multiple components). - is_multi_species: If the quantity combines several species into a single - dataset (e.g. the sound speed, which mixes the electrons and every ion). - Such a quantity is fetched once for the whole species list rather than - once per species, and its fetch function receives one list of sources per - species instead of a flat list. - """ - name = None - source = None - fetch_func = None - label = None - is_time_dep = None - is_species_dep = None - is_vector = None - is_tensor = None - is_integrated = None - is_geo = None - is_multi_species = None - - def __init__(self, name : str, source : list, fetch_func : callable, label : str, - is_time_dep : bool = False, is_species_dep : bool = False, is_vector : bool = False, - is_tensor : bool = False, is_integrated : bool = False, is_geo : bool = False, - is_multi_species : bool = False): - self.name = name - self.source = source - self.fetch_func = fetch_func - self.label = label - self.is_time_dep = is_time_dep - self.is_species_dep = is_species_dep - self.is_vector = is_vector - self.is_tensor = is_tensor - self.is_integrated = is_integrated - self.is_geo = is_geo - self.is_multi_species = is_multi_species - - # Internal methods. - - def _src_stem(self, path : str, name : str, species : str, src : str) -> str: - """ - Stem of the file name for a string source, including the trailing - separator before the frame number (geo files have no frame, so no separator). - """ - if self.is_geo: - return os.path.join(path, f"{name}-{src}") - elif self.is_species_dep: - src_ = f"{src}_" if src else "" - return os.path.join(path, f"{name}-{species}_{src_}") - else: - return os.path.join(path, f"{name}-{src}_") - - def _src_file_name(self, path : str, name : str, species : str, src : str, - frame : int | None) -> str: - """Full file name for a string source at the given frame.""" - if self.is_geo: - return f"{self._src_stem(path, name, species, src)}.gkyl" - else: - return f"{self._src_stem(path, name, species, src)}{frame}.gkyl" - - def _avail_frames_src(self, path : str, name : str, species : str, src : str, - frames : list[int] | None = None) -> set[int]: - """ - Set of available frames for a string source's file .gkyl. - Optionally restrict the search to the given list of frames. - """ - frames_avail : set[int] = set() - stem = self._src_stem(path, name, species, src) - - if frames: - candidates = (f"{stem}{f}.gkyl" for f in frames if os.path.isfile(f"{stem}{f}.gkyl")) - else: - candidates = glob.glob(f"{glob.escape(stem)}*.gkyl") - - for f in candidates: - suffix = f[len(stem):-5] - if suffix.isdigit(): - frames_avail.add(int(suffix)) - return frames_avail - - def _avail_combo_frames(self, path : str, name : str, species : str, - frames : list[int] | None = None) -> tuple[int, set[int]]: - """ - Find the first source combination whose files all exist and share the - same set of available frames. Returns (combo index, available frames). - A combination made up only of geo files is flagged with {-1}. - """ - frames_avail : set[int] = set() - combo_idx = 0 - # Check each combination of sources. - for cidx, combo in enumerate(self.source): - # Check each source for this combo. - for src in combo: - if isinstance(src, str) and self.is_geo: - # Geo files have no frame number; just check the file exists. - if not os.path.isfile(os.path.join(path, f"{name}-{src}.gkyl")): - frames_avail = set() - break - continue - - if isinstance(src, str): - frames_avail_q = self._avail_frames_src(path, name, species, src, frames) - else: - _, frames_avail_q = src._avail_combo_frames(path, name, species, frames) - - if frames_avail_q == {-1}: - # Source is a geo-only quantity: doesn't constrain frames, just needs to exist. - combo_idx = cidx - continue - - if frames_avail_q: - if not frames_avail: - frames_avail = set(frames_avail_q) - elif frames_avail_q != frames_avail: - # This source has different frames than previously checked files in - # this combo, so go to the next combo. - frames_avail = set() - break - combo_idx = cidx - else: - break - else: - # If all sources were geo files, frames_avail is still empty. - # Mark the combo as valid with {-1}. - if not frames_avail: - frames_avail = {-1} - combo_idx = cidx - - if frames_avail: - break - - return combo_idx, frames_avail - - # Public methods. - - def get_label(self, species : str | None = None, direction : str | None = None) -> str: - """Get the label for the quantity, replacing %s with species name or direction.""" - if self.is_vector: - if direction is not None: - return self.label % str(direction) - else: - return self.label % 'i' - elif self.is_species_dep: - if species is not None: - return self.label % str(species[0]) - else: - return self.label % 's' - else: - return self.label - - def get_avail_source(self, path : str, name : str, species : str, - frame_inp : str | None) -> tuple[int, list[int | None]]: - """ - Identify the source combination and list of frames needed to get this - quantity. frame_inp may be a single frame, a comma-separated list, or a - 'start:stop[:step]' range (None or ':' means all available frames). - """ - frame_list : list[int] = [] - if frame_inp is not None: - frame_inp = frame_inp.strip() - if "," in frame_inp: - frame_list = [int(f.strip()) for f in frame_inp.split(",")] - elif ":" not in frame_inp: - frame_list = [int(frame_inp)] - - # Discover available frames from any of the possible source combinations. - combo_idx, frames_avail = self._avail_combo_frames(path, name, species, frame_list) - - if not frames_avail: - raise FileNotFoundError(f"No files found for the requested quantity " - f"(path='{path}', name='{name}').") - - # Geo-only quantities have no frame number; return a single None sentinel. - if frames_avail == {-1}: - return combo_idx, [None] - - # Expand a range request against the available frames. - if len(frame_list) == 0: - frames_avail_sorted = sorted(frames_avail) - parts = frame_inp.split(":") if frame_inp else [""] - lower = int(parts[0]) if parts[0] else frames_avail_sorted[0] - upper = int(parts[1]) if len(parts) > 1 and parts[1] else frames_avail_sorted[-1] + 1 - step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 - frame_list = [f for f in frames_avail_sorted if lower <= f < upper and (f - lower) % step == 0] - - return combo_idx, frame_list - - def get_src_gdata(self, src : "str | GkQuantity", path : str, name : str, - species : str, frame : int | None, **extra) -> GData: - """ - Get the populated GData for a source, which is either a string (file - name) or a GkQuantity (computed from its own sources). - """ - if isinstance(src, str): - return GData(self._src_file_name(path, name, species, src, frame)) - - # src is a GkQuantity: resolve its own source combination and compute it. - combo_idx, _ = src.get_avail_source(path, name, species, str(frame)) - combo = src.source[combo_idx] - fetch_func = src.fetch_func[combo_idx] - gdatas = [src.get_src_gdata(s, path, name, species, frame, **extra) for s in combo] - return fetch_func(gdatas, **extra) - - def fetch(self, path : str, name : str, species : str, frame : int | None, - combo_idx : int, **extra) -> GData: - """ - Return the GData associated with this quantit by fetching the source files - and computing the quantity. - """ - combo = self.source[combo_idx] - fetch_func = self.fetch_func[combo_idx] - gdatas = [self.get_src_gdata(src, path, name, species, frame, **extra) for src in combo] - # Pass the path, name, species, and frame to the fetch function in case it needs them. - extra["path"] = path - extra["name"] = name - extra["species"] = species - extra["frame"] = frame - return fetch_func(gdatas, **extra) - - def get_avail_source_multi(self, path : str, name : str, species_list : list[str], - frame_inp : str | None) -> tuple[int, list[int | None]]: - """ - Multi-species counterpart of get_avail_source: resolve the sources for each - species in species_list and keep only the frames available for all of them, - since the quantity combines every species into one dataset. - """ - combo_idx = 0 - frames_common : set[int | None] | None = None - for species in species_list: - combo_idx, frames = self.get_avail_source(path, name, species, frame_inp) - frames_common = set(frames) if frames_common is None else frames_common & set(frames) - # end - - if not frames_common: - raise FileNotFoundError( - f"No frames are available for all of the requested species {species_list} " - f"(path='{path}', name='{name}').") - # end - - return combo_idx, sorted(frames_common, key=lambda f: (f is None, f)) - - def fetch_multi(self, path : str, name : str, species_list : list[str], - frame : int | None, combo_idx : int, **extra) -> GData: - """ - Multi-species counterpart of fetch, for is_multi_species quantities. - - The fetch function is handed one list of sources per species, in the order - of species_list: gdatas[i][j] is the j-th source of the i-th species. The - species names are passed along as extra['species']. - - Each species' sources are resolved with extra['species_idx'] set to that - species' position, so a per-species '--extra' array (e.g. 'mass=1,2,3') - picks the right entry inside the sources too, not just at the top level. - """ - combo = self.source[combo_idx] - fetch_func = self.fetch_func[combo_idx] - gdatas = [[self.get_src_gdata(src, path, name, species, frame, - **dict(extra, species_idx=species_idx)) - for src in combo] - for species_idx, species in enumerate(species_list)] - extra["path"] = path - extra["name"] = name - extra["species"] = list(species_list) - extra["frame"] = frame - return fetch_func(gdatas, **extra) - - -class GkQuantityRegistry: - """ - Registry of pre-named gyrokinetic quantities. - - Attributes: - registry: Dictionary mapping quantity names to GkQuantity objects. - """ - def __init__(self): - self.registry = {} - - def register(self, gk_quantity: GkQuantity): - """Register a new gyrokinetic quantity.""" - self.registry[gk_quantity.name] = gk_quantity - - def get(self, name: str) -> GkQuantity: - """Get a registered gyrokinetic quantity by name.""" - return self.registry.get(name) - - def list(self) -> list: - """Get a list of all registered gyrokinetic quantity names.""" - return sorted(list(self.registry.keys())) - - def has(self, name: str) -> bool: - """Check if a quantity is registered.""" - return name in self.registry diff --git a/src/postgkyl/utils/gk_quantities/registry.py b/src/postgkyl/utils/gk_quantities/registry.py deleted file mode 100644 index 4cbcd2fb..00000000 --- a/src/postgkyl/utils/gk_quantities/registry.py +++ /dev/null @@ -1,493 +0,0 @@ -""" -Registry of pre-named gyrokinetic quantities. - -Each entry is an instance of the GkQuantity class. -""" - -import postgkyl.utils.gk_quantities.fetch_funcs as ff -from .gkquantity import GkQuantity, GkQuantityRegistry - -# Instance that will hold all available gyrokinetic quantities. -gk_quant_registry: GkQuantityRegistry = GkQuantityRegistry() - -# ------------------- Register quantities ------------------- - -# ----------------------------------- -# --- Scalar geometric quantities --- -# ----------------------------------- - -# Configuration space Jacobian (interior). -_geo_int_jacobgeo : GkQuantity = GkQuantity( - name = "geo_int_jacobgeo", - source = [["geo_int_jacobgeo"],], - fetch_func = [ff.fetch_s0c0], - label = r"$J$", - is_geo = True -) -gk_quant_registry.register(_geo_int_jacobgeo) - -# Reciprocal of configuration space Jacobian (interior). -_geo_int_jacobgeo_inv : GkQuantity = GkQuantity( - name = "geo_int_jacobgeo_inv", - source = [["geo_int_jacobgeo_inv"],], - fetch_func = [ff.fetch_s0c0], - label = r"$J^{-1}$", - is_geo = True -) -gk_quant_registry.register(_geo_int_jacobgeo_inv) - -# Total Jacobian (interior). -_geo_int_jacobtot : GkQuantity = GkQuantity( - name = "geo_int_jacobtot", - source = [["geo_int_jacobtot"],], - fetch_func = [ff.fetch_s0c0], - label = r"$J$", - is_geo = True -) -gk_quant_registry.register(_geo_int_jacobtot) - -# Reciprocal of Jacobian times bmag (interior). -_geo_int_jacobtot_inv : GkQuantity = GkQuantity( - name = "geo_int_jacobtot_inv", - source = [["geo_int_jacobtot_inv"],], - fetch_func = [ff.fetch_s0c0], - label = r"$(J B)^{-1}$", - is_geo = True -) -gk_quant_registry.register(_geo_int_jacobtot_inv) - -# Magnetic field magnitude (interior). -_geo_int_bmag : GkQuantity = GkQuantity( - name = "geo_int_bmag", - source = [["geo_int_bmag"],], - fetch_func = [ff.fetch_s0c0], - label = r"$B$ (T)", - is_geo = True -) -gk_quant_registry.register(_geo_int_bmag) - -# ----------------------------------- -# --- Vector geometric quantities --- -# ----------------------------------- - -# Covariant components of magnetic field unit vector (interior). -_geo_int_b_i : GkQuantity = GkQuantity( - name = "geo_int_b_i", - source = [["geo_int_b_i"],], - fetch_func = [ff.fetch_s0cAll], - label = r"$b_%s$", - is_vector = True, - is_geo = True -) -gk_quant_registry.register(_geo_int_b_i) - -# -------------------------------------------- -# --- Field quantities (species-dependent) --- -# -------------------------------------------- - -# Electrostatic potential. -_field : GkQuantity = GkQuantity( - name = "field", - source = [["field"],], - fetch_func = [ff.fetch_s0c0], - label = r"$\phi$ (V)", - is_time_dep = True, -) -gk_quant_registry.register(_field) - -# ------------------------------------------ -# --- Plasma moments (species-dependent) --- -# ------------------------------------------ - -# Zeroth velocity moment. -_M0 : GkQuantity = GkQuantity( - name = "M0", - source = [["M0"], ["M0M1M2"], ["M0M1M2parM2perp"], ["MaxwellianMoments"], ["BiMaxwellianMoments"], ["HamiltonianMoments"],], - fetch_func = [ff.fetch_s0c0, ff.fetch_s0c0, ff.fetch_s0c0, ff.fetch_s0c0, ff.fetch_s0c0, ff.fetch_s0c0], - label = r"$M_{0%s}$ (m$^{-3}$)", - is_species_dep = True, - is_time_dep = True -) -gk_quant_registry.register(_M0) - -# First velocity moment. -_M1 : GkQuantity = GkQuantity( - name = "M1", - source = [["M1"], ["M0M1M2"], ["M0M1M2parM2perp"], ["MaxwellianMoments"], ["BiMaxwellianMoments"], ["HamiltonianMoments"],], - fetch_func = [ff.fetch_s0c0, ff.fetch_s0c1, ff.fetch_s0c1, ff.fetch_s0c0_mul_s0c1, ff.fetch_s0c0_mul_s0c1, ff.fetch_M1_from_H], - label = r"$M_{1%s}$ (m$^{-2}$/s)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M1) - -# Second parallel velocity moment. -_M2par : GkQuantity = GkQuantity( - name = "M2par", - source = [["M2par"], ["M0M1M2parM2perp"], ["M2","M2perp"]], - fetch_func = [ff.fetch_s0c0, ff.fetch_s0c2, ff.fetch_s0c0_sub_s1c0], - label = r"$M_{2\parallel%s}$ (m$^{-1}$/s$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M2par) - -# Second perpendicular velocity moment. -_M2perp : GkQuantity = GkQuantity( - name = "M2perp", - source = [["M2perp"], ["M0M1M2parM2perp"], ["M2","M2par"]], - fetch_func = [ff.fetch_s0c0, ff.fetch_s0c3, ff.fetch_s0c0_sub_s1c0], - label = r"$M_{2\perp%s}$ (m$^{-1}$/s$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M2perp) - -# Second velocity moment. -_M2 : GkQuantity = GkQuantity( - name = "M2", - source = [["M2"], ["M0M1M2"], ["M0M1M2parM2perp"], [_M2par,_M2perp],], - fetch_func = [ff.fetch_s0c0, ff.fetch_s0c2, ff.fetch_s0c2_add_s0c3, ff.fetch_s0c0_add_s1c0,], - label = r"$M_{2%s}$ (m$^{-1}$/s$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M2) - -# Third parallel velocity moment. -_M3par : GkQuantity = GkQuantity( - name = "M3par", - source = [["M3par"]], - fetch_func = [ff.fetch_s0c0], - label = r"$M_{3\parallel%s}$ (1/s$^3$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M3par) - -# Third perpendicular velocity moment. -_M3perp : GkQuantity = GkQuantity( - name = "M3perp", - source = [["M3perp"]], - fetch_func = [ff.fetch_s0c0], - label = r"$M_{3\perp%s}$ (1/s$^3$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M3perp) - -# Third velocity moment. -_M3 : GkQuantity = GkQuantity( - name = "M3", - source = [["M3"],[_M3par,_M3perp],], - fetch_func = [ff.fetch_s0c0,ff.fetch_s0c0_add_s1c0], - label = r"$M_{3%s}$ (1/s$^3$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_M3) - -# Parallel drift speed. -_upar : GkQuantity = GkQuantity( - name = "upar", - source = [["MaxwellianMoments"], ["BiMaxwellianMoments"], [_M0, _M1]], - fetch_func = [ff.fetch_s0c1, ff.fetch_s0c1, ff.fetch_s1c0_div_s0c0], - label = r"$u_{\parallel %s}$ (m/s)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_upar) - -# Parallel temperature. -_Tpar : GkQuantity = GkQuantity( - name = "Tpar", - source = [["BiMaxwellianMoments"],[_M0,_M1,_M2par],], - fetch_func = [ff.fetch_Tpar_from_BiMax, ff.fetch_Tpar_from_M0_M1_M2par], - label = r"$T_{\parallel %s}$ (J)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_Tpar) - -# Perpendicular temperature. -_Tperp : GkQuantity = GkQuantity( - name = "Tperp", - source = [["BiMaxwellianMoments"], [_M0,_M2perp]], - fetch_func = [ff.fetch_Tperp_from_BiMax, ff.fetch_Tperp_from_M0_M2perp], - label = r"$T_{\perp %s}$ (J)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_Tperp) - -# --------------------------------------------------- -# --- Combined plasma moments (species-dependent) --- -# --------------------------------------------------- - -# Temperature. -_temp : GkQuantity = GkQuantity( - name = "temp", - source = [["MaxwellianMoments"], [_Tpar,_Tperp]], - fetch_func = [ff.fetch_temp_from_Max, ff.fetch_temp_from_Tpar_Tperp], - label = r"$T_{%s}$ (J)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_temp) - -# Pressure. -_press : GkQuantity = GkQuantity( - name = "press", - source = [["MaxwellianMoments"], ["BiMaxwellianMoments"], [_M0,_temp]], - fetch_func = [ff.fetch_press_from_Max, ff.fetch_press_from_BiMax, ff.fetch_s0c0_mul_s1c0], - label = r"$p_{%s}$ (Pa)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_press) - -# Parallel pressure. -_presspar : GkQuantity = GkQuantity( - name = "presspar", - source = [[_M0,_Tpar]], - fetch_func = [ff.fetch_press_p], - label = r"$p_{\parallel %s}$ (Pa)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_presspar) - -# Perpendicular pressure. -_pressperp : GkQuantity = GkQuantity( - name = "pressperp", - source = [[_M0,_Tperp]], - fetch_func = [ff.fetch_press_p], - label = r"$p_{\perp %s}$ (Pa)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_pressperp) - -# Parallel flux of parallel energy, in the lab frame: (m/2)*M3par. -_qpar : GkQuantity = GkQuantity( - name = "qpar", - source = [[_M3par]], - fetch_func = [ff.fetch_qpar], - label = r"$q_{\parallel %s}$ (W/m$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qpar) - -# Parallel flux of perpendicular energy, in the lab frame: (m/2)*M3perp. -_qperp : GkQuantity = GkQuantity( - name = "qperp", - source = [[_M3perp]], - fetch_func = [ff.fetch_qperp], - label = r"$q_{\perp %s}$ (W/m$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qperp) - -# Parallel heat flux in the fluid (drift) frame: (m/2)*int (vpar-upar)^3 f dv. -_qpar_fluid : GkQuantity = GkQuantity( - name = "qpar_fluid", - source = [[_M0,_M1,_M2par,_M3par]], - fetch_func = [ff.fetch_qpar_fluid], - label = r"$q_{\parallel %s}^{fluid}$ (W/m$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qpar_fluid) - -# Perpendicular heat flux in the fluid (drift) frame: (m/2)*int (vpar-upar)*vperp^2 f dv. -_qperp_fluid : GkQuantity = GkQuantity( - name = "qperp_fluid", - source = [[_M0,_M1,_M2perp,_M3perp]], - fetch_func = [ff.fetch_qperp_fluid], - label = r"$q_{\perp %s}^{fluid}$ (W/m$^2$)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qperp_fluid) - -# Plasma beta. -_beta : GkQuantity = GkQuantity( - name = "beta", - source = [[_geo_int_bmag,_press],], - fetch_func = [ff.fetch_beta_from_bmag_press], - label = r"$\beta_{%s}$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_beta) - -# Thermal velocity. -_vt : GkQuantity = GkQuantity( - name = "vt", - source = [[_temp],], - fetch_func = [ff.fetch_vt], - label = r"$v_{t,%s}$ (m/s)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_vt) - -# Larmor (gyro-)radius. -_larmor_radius : GkQuantity = GkQuantity( - name = "larmor_radius", - source = [[_temp, _geo_int_bmag],], - fetch_func = [ff.fetch_larmor_radius], - label = r"$\rho_{%s}$ (m)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_larmor_radius) - -# Debye length. -_debye_length : GkQuantity = GkQuantity( - name = "debye_length", - source = [[_temp, _M0],], - fetch_func = [ff.fetch_debye_length], - label = r"$\lambda_{D,%s}$ (m)", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_debye_length) - -# Sound speed. -_c_s : GkQuantity = GkQuantity( - name = "c_s", - source = [[_M0, _temp],], - fetch_func = [ff.fetch_c_s], - label = r"$c_{s}$ (m/s)", - is_time_dep = True, - is_species_dep = False, - is_multi_species = True, -) -gk_quant_registry.register(_c_s) - -# ------------------------ -# --- Drift velocities --- -# ------------------------ - -# ExB drift velocity. -_ExB_vel : GkQuantity = GkQuantity( - name = "ExB_vel", - source = [[_geo_int_jacobtot_inv,_geo_int_bmag,_geo_int_b_i,_field],], - fetch_func = [ff.fetch_ExB_vel], - label = r"$v_{E,%s}$ (m/s)", - is_time_dep = True, - is_vector = True -) -gk_quant_registry.register(_ExB_vel) - -# Grad B drift velocity. -_gradB_vel : GkQuantity = GkQuantity( - name = "gradB_vel", - source = [[_geo_int_jacobtot_inv,_geo_int_bmag,_geo_int_b_i, _Tperp]], - fetch_func= [ff.fetch_gradB_vel], - label = r"$v_{\nabla B,%s}$ (m/s)", - is_time_dep = True, - is_species_dep = True, - is_vector = True -) -gk_quant_registry.register(_gradB_vel) - -# Diamagnetic drift velocity. -_diamag_vel : GkQuantity = GkQuantity( - name = "diamag_vel", - source = [[_geo_int_jacobtot_inv,_geo_int_bmag,_geo_int_b_i, _M0, _pressperp]], - fetch_func= [ff.fetch_diamag_vel], - label = r"$v_{dia,%s}$ (m/s)", - is_time_dep = True, - is_species_dep = True, - is_vector = True -) -gk_quant_registry.register(_diamag_vel) - -# ------------------------------ -# --- Phase space quantities --- -# ------------------------------ - -# Distribution function loaded through load_gk_distf. -_distf : GkQuantity = GkQuantity( - name = "distf", - source = [[""]], - fetch_func = [ff.load_distf], - label = r"$f_{%s}$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_distf) - -# ----------------------------- -# --- Normalized quantities --- -# ----------------------------- - -# Ratio of the Larmor radius to the Debye length. -_rho_over_lambda : GkQuantity = GkQuantity( - name = "rho_over_lambda", - source = [[_larmor_radius, _debye_length],], - fetch_func = [ff.fetch_rho_over_lambda], - label = r"$(\rho/\lambda_D)_{%s}$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_rho_over_lambda) - -# Normalized elctrostatic potential. -_phi_norm : GkQuantity = GkQuantity( - name = "phi_norm", - source = [[_field, _temp],], - fetch_func = [ff.fetch_phi_norm], - label = r"$e\phi/T_{%s}$", - is_time_dep = True, - is_species_dep = False, -) -gk_quant_registry.register(_phi_norm) - -# Normalized parallel heatflux. -_qpar_norm : GkQuantity = GkQuantity( - name = "qpar_norm", - source = [[_qpar, _M0, _temp, _vt],], - fetch_func = [ff.fetch_qpar_norm], - label = r"$q_{\parallel %s}/(n T v_{th})$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qpar_norm) - -# Normalized perpendicular heatflux. -_qperp_norm : GkQuantity = GkQuantity( - name = "qperp_norm", - source = [[_qperp, _M0, _temp, _vt],], - fetch_func = [ff.fetch_qperp_norm], - label = r"$q_{\perp %s}/(n T v_{th})$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qperp_norm) - -# Normalized parallel fluid-frame heatflux. -_qpar_fluid_norm : GkQuantity = GkQuantity( - name = "qpar_fluid_norm", - source = [[_qpar_fluid, _M0, _temp, _vt],], - fetch_func = [ff.fetch_qpar_norm], - label = r"$q_{\parallel %s}^{fluid}/(n T v_{t})$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qpar_fluid_norm) - -# Normalized perpendicular fluid-frame heatflux. -_qperp_fluid_norm : GkQuantity = GkQuantity( - name = "qperp_fluid_norm", - source = [[_qperp_fluid, _M0, _temp, _vt],], - fetch_func = [ff.fetch_qperp_norm], - label = r"$q_{\perp %s}^{fluid}/(n T v_{t})$", - is_time_dep = True, - is_species_dep = True, -) -gk_quant_registry.register(_qperp_fluid_norm) diff --git a/src/postgkyl/utils/gk_utils.py b/src/postgkyl/utils/gk_utils.py deleted file mode 100644 index 7ca16e27..00000000 --- a/src/postgkyl/utils/gk_utils.py +++ /dev/null @@ -1,291 +0,0 @@ -# -# Hardcoded parameters and auxiliary functions -# used in gyrokinetic functions. -# -import numpy as np -import os -import glob -from postgkyl.data import GInterpModal -from postgkyl.data import GData -from postgkyl.utils import verb_print -import postgkyl.utils.gkeyll_enums as gkenums - -max_num_blocks = 10000 # Maximum number of blocks. - -file_fmt = '.gkyl' # File format assumed. - -# Labels used to identify boundary flux files. -edges = ["lower","upper"] -dirs = ["x","y","z"] -# Line styles. -line_styles = ['-','--',':','-.','None','None','None','None'] -# Font sizes. -xy_label_font_size = 17 -title_font_size = 17 -tick_font_size = 14 -legend_font_size = 14 -colorbar_label_font_size = 17 - -def set_tick_font_size(axIn,fontSizeIn): - # Set the font size of the ticks to a given size. - axIn.tick_params(axis='both',labelsize=fontSizeIn) - offset_txt = axIn.yaxis.get_offset_text() # Get the text object - offset_txt.set_size(fontSizeIn) # Set the size. - offset_txt = axIn.xaxis.get_offset_text() # Get the text object - offset_txt.set_size(fontSizeIn) # Set the size. - -def read_gfile(file_name, **kwargs): - # Read a Gkeyll file. - c2p_file = None - if "mapc2p" in kwargs.keys(): - c2p_file = kwargs["mapc2p"] - - pgData = GData(file_name, mapc2p_name=c2p_file) # Read data with pgkyl. - grid = pgData.get_grid() # Time stamps of the simulation. - vals = pgData.get_values() # Data values. - if isinstance(grid, np.ndarray): - grid_out = np.squeeze(grid) - else: - grid_out = list() - for d in range(len(grid)): - grid_out.append(np.squeeze(grid[d])) - - return grid_out, np.squeeze(vals), pgData - -def read_gfile_if_present(file_name, **kwargs): - # Check if a Gkeyll file exists. If it does, read it and return - # its grid, data and GData object. If it doesn't, return None. - if os.path.exists(file_name): - grid, vals, pgdat = read_gfile(file_name, kwargs) - return True, np.squeeze(grid), np.squeeze(vals), pgdat - else: - verb_print(ctx, " -> File "+file_name+" not found. Proceeding w/o it.") - return False, None, None, None - -def read_interp_gfile(file_name, poly_order, basis_type, comp=0): - # Read a Gkeyll file and interpolate its DG dataset assuming it has a - # polynomial basis of 'poly_order' order and basis type 'basis_type'. - # Optional argument 'comp' requests a specific component if a file - # contains multiple DG datasets. - pgData = GData(file_name) # Read data with pgkyl. - interp = GInterpModal(pgData,poly_order,basis_type) - grid, vals = interp.interpolate(comp) - if isinstance(grid, np.ndarray): - grid_out = np.squeeze(grid) - else: - grid_out = list() - for d in range(len(grid)): - grid_out.append(np.squeeze(grid[d])) - - return grid_out, np.squeeze(vals), pgData - -def dict_get_bool(dict_in, key, default): - # Interpret a dictionary value as a bool, returning 'default' if the key is - # absent. String values '1'/'true' (case-insensitive) are True, anything - # else false. Non string values are converted using bool(). - if key not in dict_in: - return default - val = dict_in[key] - if isinstance(val, str): - return val.strip().lower() in ("1", "true") - return bool(val) - -def parse_slice_string(value): - # Parse a 'slice()' from string, like 'start:stop:step'. - parts = value.split(':') - # Convert parts to integers, replacing empty strings with None for slice defaults - parsed_parts = [] - for p in parts: - try: - parsed_parts.append(int(p) if p else None) - except ValueError: - # Handle cases where the part might not be a number - raise ValueError(f"Invalid slice part: {p}") - # Create the slice object with the appropriate number of arguments - return slice(*parsed_parts) - -def get_block_indices(multib, file_path_name): - # Return a list of the indices of the blocks in a multiblock simulation - # to be processed. - # - multib: ="-10" single block. - # ="-1" will find all the blocks. - # =comma-separated list or slice of desired blocks to use. - # - file_path_name: path and file name used to find blocks, with block - # index replaced by "*", e.g. "_b*-_field_0.gkyl". - def is_str_an_int(str_in): - try: - int(str_in) - return True - except ValueError: - return False - # end - # end - - if multib == "-10": - # Single block. - blocks = [0] - else: - # Multi block. - if multib == "-1": - # Find and use all blocks. - file_list = glob.glob(file_path_name) - num_blocks = len(file_list) - blocks = list(range(num_blocks)) - else: - # Use specified blocks. - if ',' in multib: - blocks = multib.split(",") - num_blocks = len(blocks) - blocks = [int(blocks[i]) for i in range(num_blocks)] - elif ':' in multib: - slice_obj = parse_slice_string(multib) - blocks = list(range(*slice_obj.indices(max_num_blocks))) - elif is_str_an_int(multib): - blocks = [int(multib)] - else: - raise NameError("Blocks given to --multib -m must be a comma separated list or slice.") - - return blocks - -def nodes_to_RZ(nodes, is_mapc2p): - """Extract 2D R and Z node arrays from a nodes.gkyl data array. - - For 3D config space, the phi=0 slice is used. - """ - nx_nod = np.shape(nodes) - cdim = np.size(nx_nod) - 1 - cart_dim = 3 - lo_idx = [[0] * cdim + [cd] for cd in range(cart_dim)] - up_idx = [[nx_nod[d] for d in range(cdim)] + [cd + 1] for cd in range(cart_dim)] - - if cdim == 3: - yidx = 0 - for cd in range(cart_dim): - lo_idx[cd][1] = yidx - up_idx[cd][1] = yidx + 1 - # end - - slices = [[slice(lo_idx[cd][d], up_idx[cd][d]) for d in range(cdim + 1)] for cd in range(cart_dim)] - - if is_mapc2p: - cartX = [np.squeeze(nodes[tuple(slices[d])]) for d in range(cart_dim)] - majorR = np.sqrt(cartX[0] ** 2 + cartX[1] ** 2) - vertZ = cartX[2] - else: - majorR = np.squeeze(nodes[tuple(slices[0])]) - vertZ = np.squeeze(nodes[tuple(slices[1])]) - # end - - return majorR, vertZ - - -def is_gdata_geo_mapc2p(gdata): - # Determine whether the GData object, gdata, is from a simulation with MAPC2P - # geometry. If geometry_type is missing from the metadata, default to true. - gdata_meta = gdata.get_ctx() - is_mapc2p = True - if ("geometry_type" in gdata_meta): - if "geometry_type" in gdata_meta.keys(): - mc2p_idx = gkenums.enum_key_to_idx(gkenums.gkyl_geometry_id,"GKYL_GEOMETRY_MAPC2P") - is_mapc2p = mc2p_idx == gdata_meta["geometry_type"] - # end - #end - return is_mapc2p - -#quant_attributes = { -# "den" : { -# "files" : [["M0"],["MaxwellianMoments"],["BiMaxwellianMoments"],], -# "fetch_func" : [fccq_get_c0, fccq_get_c0, fccq_get_c0,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "norm_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$n_%s$ (m$^{-3}$)' -# }, -# "upar" : { -# "files" : [["MaxwellianMoments"],["BiMaxwellianMoments"],["M0","M1"],], -# "fetch_func" : [fccq_get_c1, fccq_get_c1, fccq_get_f1Df0,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$u_{\parallel %s}$ (m$^{-3}$)' -# }, -# "tpar" : { -# "files" : [["BiMaxwellianMoments"],], -# "fetch_func" : [fccq_get_c2,], -# "scale_func" : [fccq_scale_mDe,], -# "label" : r'$T_{\parallel %s}$ (eV)' -# }, -# "tperp" : { -# "files" : [["BiMaxwellianMoments"],], -# "fetch_func" : [fccq_get_c3, ], -# "scale_func" : [fccq_scale_mDe,], -# "label" : r'$T_{\perp %s}$ (eV)' -# }, -# "temp" : { -# "files" : [["MaxwellianMoments"],["BiMaxwellianMoments"],], -# "fetch_func" : [fccq_get_c2, fccq_2c3Pc2D3], -# "scale_func" : [fccq_scale_mDe,], -# "label" : r'$T_{%s}$ (eV)' -# }, -# "m0" : { -# "files" : [["M0"],["M0M1M2"],["M0M1M2PARM2PERP"],["MaxwellianMoments"],["BiMaxwellianMoments"],], -# "fetch_func" : [fccq_get_c0, fccq_get_c0, fccq_get_c0, fccq_get_c0, fccq_get_c0,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$M_{0 %s}$ (m$^{-3}$)' -# }, -# "m1" : { -# "files" : [["M1"],["M0M1M2"],["M0M1M2PARM2PERP"],], -# "fetch_func" : [fccq_get_c0, fccq_get_c1, fccq_get_c1,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$M_{1 %s}$ (m$^{-2}$ s^{-1})' -# }, -# "m2par" : { -# "files" : [["M2PAR"],["M0M1M2PARM2PERP"],], -# "fetch_func" : [fccq_get_c0, fccq_get_c2,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$M_{2\parallel %s}$ (m$^{-1}$ s^{-2})' -# }, -# "m2perp" : { -# "files" : [["M2PERP"],["M0M1M2PARM2PERP"],], -# "fetch_func" : [fccq_get_c0, fccq_get_c3,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$M_{2\perp %s}$ (m$^{-1}$ s^{-2})' -# }, -# "m2" : { -# "files" : [["M2"],["M0M1M2"],["M0M1M2PARM2PERP"],["M2PAR","M2PERP"],], -# "fetch_func" : [fccq_get_c0, fccq_get_c2, fccq_get_c2Pc3, fccq_get_f0Pf1,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$M_{2 %s}$ (m$^{-1}$ s^{-2})' -# }, -# "m3par" : { -# "files" : [["M3PAR"],], -# "fetch_func" : [fccq_get_c0,], -# "scale_func" : [fccq_scale_disabled,], -# "label" : r'$M_{3\parallel %s}$ (s^{-3})' -# }, -# "m3perp" : { -# "files" : [["M3PERP"],], -# "fetch_func" : [fccq_get_c0,], -# "scale_func" : [fccq_scale_disabled,], -# "label" : r'$M_{3\perp %s}$ (s^{-3})' -# }, -# "m3" : { -# "files" : [["M3"],["M3PAR","M3PERP"],], -# "fetch_func" : [fccq_get_c0, fccq_get_f0Pf1,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$M_{3 %s}$ (s^{-3})' -# }, -# "phi" : { -# "files" : [["field"],], -# "fetch_func" : [fccq_get_c0,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$\phi$ (V)' -# }, -# "bmag" : { -# "files" : [["bmag"],], -# "fetch_func" : [fccq_get_c0,], -# "scale_func" : [fccq_scale_disabled, fccq_scale_disabled, fccq_scale_disabled,], -# "label" : r'$B$ (T)' -# }, -#} - -# -# End of hardcoded parameters and auxiliary functions. -# diff --git a/src/postgkyl/utils/gkeyll_const.py b/src/postgkyl/utils/gkeyll_const.py deleted file mode 100644 index 3233d818..00000000 --- a/src/postgkyl/utils/gkeyll_const.py +++ /dev/null @@ -1,14 +0,0 @@ -# Universal constants. Maybe we can just load these from the gkeyll library. - -GKYL_PI = 3.141592653589793238462643383279502884 -GKYL_E = 2.718281828459045235360287471352662497 -GKYL_SPEED_OF_LIGHT = 299792458.0 # m/s -GKYL_PLANCKS_CONSTANT_H = 6.62606896e-34 # joule*seconds -GKYL_ELECTRON_MASS = 9.10938215e-31 # Kg -GKYL_PROTON_MASS = 1.672621637e-27 # Kg -GKYL_MASS_UNIT = 1.66053907e-27 # Kg -GKYL_ELEMENTARY_CHARGE = 1.602176487e-19 # Coulomb -GKYL_BOLTZMANN_CONSTANT = 1.3806488e-23 -GKYL_EPSILON0 = 8.854187817620389850536563031710750260608e-12 # farad/meter -GKYL_MU0 = 12.56637061435917295385057353311801153679e-7 # newtons/ampere/ampere -GKYL_EV2KELVIN = GKYL_ELEMENTARY_CHARGE/GKYL_BOLTZMANN_CONSTANT diff --git a/src/postgkyl/utils/gkeyll_enums.py b/src/postgkyl/utils/gkeyll_enums.py deleted file mode 100644 index fe127d4c..00000000 --- a/src/postgkyl/utils/gkeyll_enums.py +++ /dev/null @@ -1,47 +0,0 @@ -# -# A set of enums in gkeyll. They have to match those in the Gkeyll source code. -# - -# Identifiers for specific geometry types -gkyl_geometry_id = [ - "GKYL_GEOMETRY_NONE", # No geometry, use Cartesian. - "GKYL_GEOMETRY_TOKAMAK", # Tokamak Geometry from Efit. - "GKYL_GEOMETRY_MIRROR", # Mirror Geometry from Efit. - "GKYL_GEOMETRY_MAPC2P", # General geometry from user provided mapc2p. - "GKYL_GEOMETRY_FROMFILE", # Geometry from file. -] - -gkyl_basis_type = [ - "GKYL_BASIS_MODAL_SERENDIPITY", - "GKYL_BASIS_MODAL_TENSOR", - "GKYL_BASIS_MODAL_HYBRID", - "GKYL_BASIS_MODAL_GKHYBRID", - "GKYL_BASIS_MODAL_GKHYBRID_VEL", -] - -pgkyl_basis_type = [ - "serendipity", - "tensor", - "hybrid", - "gkhybrid", - "gkhybrid_vel", -] - -def enum_idx_to_key(enum, idx): - # Given an enum list, return the string corresponding to the index idx - # provided. - return enum[idx]; - -def enum_key_to_idx(enum, key): - # Given an enum list, return the index of the string key provided. - return enum.index(key); - -def basis_type_gkyl_to_pgkyl(gkyl_basis_type_in): - # Convert the basis type given as a gkeyll enum int or string, - # to the string used the rest of postgkyl. - if isinstance(gkyl_basis_type_in, int): - return pgkyl_basis_type[gkyl_basis_type_in] - elif isinstance(gkyl_basis_type_in, str): - return pgkyl_basis_type[enum_key_to_idx(gkyl_basis_type,gkyl_basis_type_in)] - else: - ValueError("Wrong input to basis_type_gkyl_to_pgkyl.") diff --git a/src/postgkyl/utils/input_parser.py b/src/postgkyl/utils/input_parser.py deleted file mode 100644 index 7ce26468..00000000 --- a/src/postgkyl/utils/input_parser.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Postgkyl module to unify inputs for various tools and diagnostics.""" -from __future__ import annotations - -import numpy as np -from typing import Tuple, TYPE_CHECKING - -if TYPE_CHECKING: - from postgkyl import GData -# end -import postgkyl.data.gdata - -def input_parser(data: GData | np.ndarray | Tuple[list, np.ndarray]) -> Tuple[list, np.ndarray]: - """Utility function to parse input and return grid and values. - - Motivation for this funtion is to unify what input is used by Postgkyl tools and - diagnostics. Sometimes it's beneficial to pass the internal GData class and in other - situations it's more convenient to pass a grid and values. - - Args: - data: GData | NumPy array | tuple of grid list and NumPy array - Input ot be parsef - - Returns: - grid: list of NumPy arrays - values: NumPy array - - Raises: - TypeError when wrong data type is provided - ValueError dimensions of grid and values don't match - """ - if isinstance(data, postgkyl.data.gdata.GData): - return data.get_grid(), data.get_values() - elif isinstance(data, np.ndarray): - return (), data - elif isinstance(data, tuple) or isinstance(data, list): # A little leeway - if len(data) == 2: - if not isinstance(data[0], list): - raise TypeError("Input grid needs to be a list of NumPy arrays.") - if not isinstance(data[1], np.ndarray): - raise TypeError("Input values needs to be a NumPy array.") - if len(data[0]) != len(data[1].shape) and len(data[0]) != len(data[1].shape)-1: - raise ValueError("Input grid and valeus don't have the same number of dimesnions.") - return data[0], data[1] - else: - raise TypeError("Input tuple needs to have two components: grid and values; {len(data):d} were provided.") - else: - raise TypeError("Input must be either GData class or a tuple of grid and values.") - # end diff --git a/src/postgkyl/utils/load_style.py b/src/postgkyl/utils/load_style.py deleted file mode 100644 index 2c2d51d1..00000000 --- a/src/postgkyl/utils/load_style.py +++ /dev/null @@ -1,17 +0,0 @@ -from cycler import cycler -import click - -def load_style(ctx: click.core.Context, fn: str) -> None: - fh = open(fn, "r", encoding="utf-8") - for line in fh.readlines(): - key = line.split(":")[0] - key_len = int(len(key)) - key = key.strip() - value = line[(key_len + 1) :].strip() - if value[:6] == "cycler": - arg = eval(value[16:-1]) - value = cycler(color=arg) - # end - ctx.obj["rcParams"][key] = value - # end - fh.close() diff --git a/src/postgkyl/utils/set_frame.py b/src/postgkyl/utils/set_frame.py deleted file mode 100644 index 3fd8b13b..00000000 --- a/src/postgkyl/utils/set_frame.py +++ /dev/null @@ -1,57 +0,0 @@ -import numpy as np -import click - -#sets frame in block ctx attribute using block file name -def set_frame(ctx: click.core.Context) -> list: - """Utility function which sets data ctx frames in multiblock data situations - - This function uses gkyl's default file name output in multiblock cases to - identify the respective frame for each loaded in data object. It assigns the correct - frame to each data object's ctx frame attribute. It then returns a list with all the - identified frames in ascending order. - - The motivation for this function is to allow for easy organization of multiblock data - objects in plotting and animation. - - Args: - ctx: click.core.context | Object - Context from loaded data / previous commands - Returns: - sorted_frame_list: list - """ - - data = ctx.obj["data"] - - #load in file names - files = [dat._file_name for dat in data.iterator()] - - #iterate through file names and find smallest index where file names differ, this is where the file name is - #this is assuming that the file names are default from gkyl - #short file is used to iterate in order to prevent indexing error - short_file = min(files, key=len) - num_frame_idx = np.inf - for i in range(len(files)): - for j in range(len(short_file)): - if short_file[j] != files[i][j] and j < num_frame_idx: - num_frame_idx = j - #end - #end - #end - - #isolate frame number in file name and append it to big frame_list - frame_list = [] - for f in files: - f = f.split(".gkyl")[0] - frame = f[num_frame_idx:].split("_")[0] - frame_list.append(int(frame)) - #end - - #data objects in iterator have same index as corresponding frame in frame_list - #this loop sets frame ctx attribute - for i, dat in data.iterator(enum=True): - dat.ctx["frame"] = frame_list[i] - #end - - #return sorted frame list for use in animate function - sorted_frame_list = np.unique(np.sort(frame_list)) - return sorted_frame_list diff --git a/src/postgkyl/utils/verb_print.py b/src/postgkyl/utils/verb_print.py deleted file mode 100644 index c418798c..00000000 --- a/src/postgkyl/utils/verb_print.py +++ /dev/null @@ -1,8 +0,0 @@ -from time import time -import click - -def verb_print(ctx: click.core.Context, message: str) -> None: - if ctx.obj["verbose"]: - elapsed_time = time() - ctx.obj["start_time"] - click.echo(click.style(f"[{elapsed_time:f}] {message:s}", fg="green")) - # end diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..dd02ab75 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,190 @@ +"""Shared pytest configuration for the postgkyl test suite. + +No-GUI guarantee +----------------- +This test session must never put a window or a browser tab on the desktop -- +doing so can crash/hang the sandboxed environment this suite runs in. Two +independent guards enforce that, both applied before any test module (or its +imports) can run: + +- ``matplotlib.use("Agg")`` is forced at import time, below, before anything + else gets a chance to trigger Matplotlib's own backend auto-detection + (which, given a display, could pick an interactive GUI backend). Agg is a + pure-raster, no-window backend, so ``plt.show()`` is always a no-op under it. +- ``_block_gui_popups`` (autouse, session-scoped) monkeypatches + ``webbrowser.open`` -- the mechanism ``render.plotly``'s ``open_preview`` + (default-off; see its docstring) uses to show a Plotly figure -- and, if + PyVista is installed, ``pyvista.Plotter.show`` -- the analogous mechanism + for a PyVista render window (see ``render.pyvista.pyvista``'s ``no_show`` + parameter). Plotly defaults ``show`` to ``False``; interactive-by-default + PyVista calls in tests pass ``no_show=True``. This fixture is the backstop + for any call that forgets to select its renderer's headless setting. + +Test data generation +-------------------- +``pytest_configure`` writes synthetic .gkyl files to +``tests/test_data/generated/`` (gitignored -- every test that reads from +that directory depends on this running first). It is a hook, not a +session-scoped autouse fixture, specifically so it runs exactly once in the +true parent process before collection or any forking begins -- see its own +comment for why a fixture is the wrong tool once macOS CI's --forked is in +play. Without it, a clean checkout (e.g. CI) has no fixtures to read; only a +machine where someone has run ``python tests/generate_test_data.py`` (or a +prior pytest session) before would happen to have them already on disk. +""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import pytest + +from generate_test_data import generate_all + +GEN_DIR = Path(__file__).parent / "test_data" / "generated" +_COMPATIBILITY_TEST_FILES = frozenset({ + "test_cli_generator.py", + "test_diagnostics_discovery.py", + "test_io_mapping.py", +}) + + +def _env_enabled(name: str) -> bool: + """Read a CI feature switch without treating ``"0"`` as truthy.""" + value = os.environ.get(name, "").strip().lower() + if not value: + return False + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + raise pytest.UsageError( + f"{name} must be one of 1/0, true/false, yes/no, or on/off; got " + f"{os.environ[name]!r}") + + +def _require_gkeyll_when_requested() -> None: + """Turn a missing native capability into a CI failure, not mass skips.""" + if not _env_enabled("POSTGKYL_REQUIRE_GKEYLL"): + return + + from postgkyl import gpython + if not gpython.available(): + pytest.exit( + "POSTGKYL_REQUIRE_GKEYLL is enabled, but the compiled Gkeyll/gpython " + "capability is unavailable. Native tests would otherwise be silently " + "skipped.", + returncode=2) + + +# macOS-only escape hatch: ``postgkyl``'s facade (__init__.py -> render -> +# render.pyvista) unconditionally imports PyVista's VTK bindings, so every +# test process here loads the full native VTK stack regardless of whether +# any PyVista test runs. On macOS CI (observed on Python 3.10, 3.11, and +# 3.12 alike -- not Python-version-specific, so not the intermittent +# matplotlib font-rendering issue tracked elsewhere in this suite) VTK's own +# C++ global/static destructors reproducibly SIGSEGV during CPython's +# interpreter finalization, always *after* pytest has already run every +# test and printed its full (passing) summary. Nothing of value happens in +# that teardown window, so once pytest has reported its result, exit the +# process immediately via ``os._exit`` -- bypassing the interpreter +# finalization that walks into VTK's broken destructors -- instead of +# letting a native-library crash overwrite an already-successful result +# with a misleading CI failure. +_exit_status: list[int] = [] + + +def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: + _exit_status.append(int(exitstatus)) + + +def pytest_unconfigure(config: pytest.Config) -> None: + if sys.platform == "darwin" and _exit_status: + sys.stdout.flush() + sys.stderr.flush() + os._exit(_exit_status[0]) + + +@pytest.fixture(scope="session", autouse=True) +def _block_gui_popups(): + """Backstop: no test may open a browser tab or a native render window.""" + import webbrowser + + def _no_browser(*_args, **_kwargs): + raise AssertionError( + "webbrowser.open() was called during tests -- a figure/preview would " + "have popped up on the desktop. Pass show=False for render.plotly or " + "no_show=True for render.pyvista, or mock the call being tested.") + + webbrowser.open = _no_browser + webbrowser.open_new = _no_browser + webbrowser.open_new_tab = _no_browser + + try: + import pyvista + + def _no_plotter_show(*_args, **_kwargs): + raise AssertionError( + "pyvista.Plotter.show() was called during tests -- a render window " + "would have popped up on the desktop. Pass no_show=True or mock " + "the call being tested.") + + pyvista.Plotter.show = _no_plotter_show + except ImportError: + pass + + +@pytest.fixture(autouse=True) +def _close_matplotlib_figures(): + """Release every pyplot-managed figure after each test.""" + yield + + import matplotlib.pyplot as plt + plt.close("all") + + +def pytest_configure(config: pytest.Config) -> None: + _require_gkeyll_when_requested() + + # A session-scoped autouse *fixture* only actually runs on first request, + # which lands inside whichever test forks first under macOS CI's --forked + # (see test.yml) -- pytest's "already cached" bookkeeping then lives in + # that child's forked copy of the session, never propagating back to the + # parent, so every subsequent forked test re-triggers it too (measured: + # 40 re-generations for 40 tests instead of one). pytest_configure runs + # exactly once, in the true parent process, before collection or any + # forking begins, so this is immune to that regardless of --forked. + generate_all(GEN_DIR) + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + """Attach capability categories from their authoritative test metadata. + + Native tests already declare a skip condition whose reason names the + compiled Gkeyll capability. Deriving ``native`` from that marker avoids a + second hand-maintained inventory. The compatibility inventory names only + modules proven to pass with the extension absent; all numerics modules are + included mechanically because that layer has no internal imports. Render + modules have one job, so their filename is likewise the single category + rule; external-tool and slow markers stay explicit on the individual tests + that actually cross a process boundary or take appreciable time. + """ + for item in items: + file_name = Path(str(item.path)).name + if (file_name.startswith("test_numerics_") + or file_name in _COMPATIBILITY_TEST_FILES): + item.add_marker("compatibility") + + if file_name.startswith("test_render_"): + item.add_marker("render") + + for marker in item.iter_markers(name="skipif"): + reason = str(marker.kwargs.get("reason", "")).lower() + if "compiled gkeyll" in reason: + item.add_marker("native") + break diff --git a/tests/generate_test_data.py b/tests/generate_test_data.py new file mode 100644 index 00000000..d84786f2 --- /dev/null +++ b/tests/generate_test_data.py @@ -0,0 +1,483 @@ +"""Generate synthetic .gkyl test files for the postgkyl test suite. + +Run directly to regenerate: + python tests/generate_test_data.py + +Called automatically by conftest.py at the start of each pytest session. + +Field files encode polyOrder and basisType in their msgpack metadata block so +GData auto-populates ctx["poly_order"] and ctx["basis_type"] on load. + +C2P mapping files store modal DG coefficients for analytical coordinate +transformations. The basis is inferred by GData from num_comps/ndim via +_get_basis_p(). Two mapping types are provided: + - "stretch": linear map (comp domain [0,1]^n → physical domain phys_bounds) + - "rotation": 2D rotation by angle α about the origin + +Dynvector files (file_type=2, a bare time series with no spatial grid -- +e.g. a field-energy history) are written by ``write_gkyl_dynvector``. +""" +import struct +from pathlib import Path + +import msgpack +import numpy as np + +_RNG = np.random.default_rng(42) +_SQRT3 = np.sqrt(3) + +# Component counts per basis -- mirrors the tables in src/postgkyl/data/dg.py +# serendipity: indexed as [ndim-1][poly_order] (p=0 → 1 component) +_COMPS_SER = [ + [1, 2, 3, 4, 5], # 1D + [1, 4, 8, 12, 17], # 2D + [1, 8, 20, 32, 50], # 3D +] +# tensor: indexed as [ndim-1][poly_order-1] (p starts at 1) +_COMPS_TEN = [ + [2, 3, 4, 5], # 1D + [4, 9, 16, 25], # 2D + [8, 27, 64, 125], # 3D +] +# maximal-order: indexed as [ndim-1][poly_order-1] +_COMPS_MAX = [ + [2, 3, 4, 5], # 1D + [3, 6, 10, 15], # 2D + [4, 10, 20, 35], # 3D +] + +_COMPS = { + "serendipity": (_COMPS_SER, lambda p: p), + "tensor": (_COMPS_TEN, lambda p: p - 1), + "maximal-order": (_COMPS_MAX, lambda p: p - 1), +} + + +def num_comps(basis: str, ndim: int, poly_order: int) -> int: + table, idx_fn = _COMPS[basis] + return table[ndim - 1][idx_fn(poly_order)] + + +def write_gkyl_field( + path: Path, + cells: list[int], + lower: list[float], + upper: list[float], + values: np.ndarray, + poly_order: int, + basis_type: str, + time: float = 0.0, + frame: int = 0, +) -> None: + """Write a minimal valid .gkyl v1 binary field file with msgpack metadata.""" + ndim = len(cells) + nc = values.shape[-1] + + meta = msgpack.packb({ + "polyOrder": poly_order, + "basisType": basis_type, + "time": time, + "frame": frame, + }) + + with open(path, "wb") as f: + # --- version-1 header --- + f.write(b"gkyl0") + f.write(struct.pack(" None: + """Write a minimal valid .gkyl v1 binary dynvector (file_type=2) file. + + A dynvector has no spatial grid -- just a time series of ``num_comps`` + values per sample (e.g. a field-energy history). Layout per + ``gkyl_reader.py``'s ``_read_t2_v1``: header, real_type, esznc, size, + then all of TIME_DATA followed by all of DATA (C-order). + """ + nc = values.shape[-1] + size = len(time) + + with open(path, "wb") as f: + f.write(b"gkyl0") + f.write(struct.pack(" float64 + esznc = nc * 8 # element_size * num_comps (bytes) + f.write(struct.pack(" np.ndarray: + """Modal DG coefficients for a linear stretch mapping (comp [0,1]^n → phys). + + For each cell the mapping is: + coord_d(xi') = coord_mid_d + (dx_phys_d/2) * xi'_d + Modal serendipity coefficients (any poly_order): + c_0 = 2 * coord_mid (constant mode, normalized by 1/2) + c_{d+1} = dx_phys / sqrt(3) (linear mode in direction d) + all higher modes = 0 (linear function has no quadratic terms) + """ + ndim = len(cells) + dx = [(phys_hi[d] - phys_lo[d]) / cells[d] for d in range(ndim)] + values = np.zeros((*cells, ndim * num_modes)) + + grids = np.meshgrid(*[np.arange(cells[d]) for d in range(ndim)], + indexing="ij") + for d in range(ndim): + mid = phys_lo[d] + dx[d] * (grids[d] + 0.5) + off = d * num_modes + values[..., off] = 2.0 * mid # constant mode + values[..., off + 1 + d] = dx[d] / _SQRT3 # linear mode in d-th direction + return values + + +def _c2p_rotation_values( + cells: list[int], + comp_lo: list[float], + comp_hi: list[float], + angle: float, + num_modes: int, +) -> np.ndarray: + """Modal DG coefficients for a 2D rotation mapping by *angle* radians. + + The computational domain is [comp_lo[0], comp_hi[0]] x [comp_lo[1], comp_hi[1]]. + The mapping is x = xi*cos - eta*sin, y = xi*sin + eta*cos. + Only valid for 2D serendipity; the linear rotation is exact at any poly_order. + """ + assert len(cells) == 2, "rotation mapping only implemented for 2D" + ca, sa = np.cos(angle), np.sin(angle) + N_x, N_y = cells + dxi = (comp_hi[0] - comp_lo[0]) / N_x + deta = (comp_hi[1] - comp_lo[1]) / N_y + + ii, jj = np.mgrid[0:N_x, 0:N_y] + xi_mid = comp_lo[0] + dxi * (ii + 0.5) + eta_mid = comp_lo[1] + deta * (jj + 0.5) + + x_mid = xi_mid * ca - eta_mid * sa + y_mid = xi_mid * sa + eta_mid * ca + + values = np.zeros((N_x, N_y, 2 * num_modes)) + # x-coordinate modal coefficients + values[..., 0] = 2.0 * x_mid # constant + values[..., 1] = dxi * ca / _SQRT3 # xi'-mode (dx/dxi' * mapping factor) + values[..., 2] = -deta * sa / _SQRT3 # eta'-mode + # y-coordinate modal coefficients + values[..., num_modes + 0] = 2.0 * y_mid + values[..., num_modes + 1] = dxi * sa / _SQRT3 + values[..., num_modes + 2] = deta * ca / _SQRT3 + return values + + +# --------------------------------------------------------------------------- +# Analytic four-component 1-D profile +# --------------------------------------------------------------------------- + + +def _mirror_comparison_profiles(z: np.ndarray, alpha: float) -> np.ndarray: + """Symmetric beam profiles used by the alpha-convergence example. + + Density and both temperatures are even in ``z``; parallel velocity is + odd. ``alpha`` introduces a small, pedestal-localized difference so the + two generated datasets remain close enough to read as a convergence + comparison. Components are already in the plotting units used by the + example: m^-3, m/s, keV, and keV. + """ + if alpha <= 0.0: + raise ValueError("alpha must be positive") + + radius = np.abs(z) + sensitivity = np.log10(alpha / 2.0e-5) + pedestal = np.exp(-((radius - 0.86) / 0.16)**2) + + density = (1.01e13 + 2.45e19 / (1.0 + np.exp( + (radius - 0.92) / 0.08)) + 4.5e18 * np.exp(-(radius / 0.55)**4)) + velocity = (1.30e6 * np.tanh(z / 0.05) / (1.0 + np.exp( + (0.82 - radius) / 0.05))) + t_parallel = (0.101 + 11.8 / (1.0 + np.exp( + (radius - 0.90) / 0.10)) + 0.4 * np.exp(-((radius - 0.55) / 0.15)**2)) + t_perpendicular = (0.101 + 20.8 / (1.0 + np.exp( + (radius - 0.98) / 0.13)) + 10.0 * np.exp(-((radius - 0.86) / 0.10)**2)) + + return np.stack([ + density * (1.0 - 0.025 * sensitivity * pedestal), + velocity * (1.0 + 0.020 * sensitivity * pedestal), + t_parallel * (1.0 + 0.035 * sensitivity * pedestal), + t_perpendicular * (1.0 + 0.030 * sensitivity * pedestal), + ], + axis=-1) + + +def _project_1d_p1(fn, lower: float, upper: float, cells: int, + *args) -> np.ndarray: + """Cellwise L2 projection of a vector-valued analytic function onto p1. + + The 1-D modal serendipity basis is ``(1/sqrt(2), sqrt(3/2)*xi)`` on + ``[-1, 1]``. Eight-point Gauss quadrature resolves the smooth analytic + profiles well within each cell. The result is field-blocked as Gkeyll + expects: ``[f0_mode0, f0_mode1, f1_mode0, ...]``. + """ + xi, weights = np.polynomial.legendre.leggauss(8) + dz = (upper - lower) / cells + centers = lower + (np.arange(cells) + 0.5) * dz + z = centers[:, None] + 0.5 * dz * xi[None, :] + samples = fn(z, *args) + basis = np.stack([ + np.full_like(xi, 1.0 / np.sqrt(2.0)), + np.sqrt(3.0 / 2.0) * xi, + ], + axis=-1) + coefficients = np.einsum("cqf,qb,q->cfb", samples, basis, weights) + return coefficients.reshape(cells, -1) + + +# --------------------------------------------------------------------------- +# Configuration tables +# --------------------------------------------------------------------------- + +# (stem, ndim, cells, poly_order, basis_type) +_FIELD_CONFIGS: list[tuple] = [ + ("1d_ms_p1", 1, [8], 1, "serendipity"), + ("1d_ms_p2", 1, [8], 2, "serendipity"), + ("2d_ms_p1", 2, [8, 8], 1, "serendipity"), + ("2d_ms_p2", 2, [8, 8], 2, "serendipity"), + ("2d_mt_p1", 2, [8, 8], 1, "tensor"), + ("2d_mt_p2", 2, [8, 8], 2, "tensor"), + ("2d_mo_p1", 2, [8, 8], 1, "maximal-order"), + ("2d_mo_p2", 2, [8, 8], 2, "maximal-order"), + ("3d_ms_p1", 3, [4, 4, 4], 1, "serendipity"), +] + +# C2P mapping files. +# (stem, kind, cells, poly_order, basis_type, extra...) +# kind="stretch": extra = (phys_lo, phys_hi) -- comp domain [0,1]^n +# kind="rotation": extra = (angle,) -- comp domain [0,1]^2 +_C2P_CONFIGS: list[tuple] = [ + # Linear stretch: physical x∈[0,2], y∈[0,3]; paired with 2d_ms_p1.gkyl + ("2d_c2p_stretch_ms_p1", "stretch", [8, 8], 1, "serendipity", [0.0, 0.0], + [2.0, 3.0]), + # Same stretch for p=2; paired with 2d_ms_p2.gkyl + ("2d_c2p_stretch_ms_p2", "stretch", [8, 8], 2, "serendipity", [0.0, 0.0], + [2.0, 3.0]), + # Rotation by 45°; paired with 2d_ms_p1.gkyl (comp domain [0,1]^2) + ("2d_c2p_rot45_ms_p1", "rotation", [8, 8], 1, "serendipity", np.pi / 4), +] + + +def generate_all(out_dir: Path | str) -> None: + """Write all synthetic test files to *out_dir*.""" + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # Stationary shock-tube initial state: p0 modal coefficients are sqrt(2) + # times the physical values in the orthonormal 1D basis. + x = (np.arange(100) + 0.5) / 100 + rho = np.where(x < 0.5, 1.0, 0.125) + pressure = np.where(x < 0.5, 1.0, 0.1) + moments = np.stack([ + rho, + np.zeros_like(x), + np.zeros_like(x), + np.zeros_like(x), pressure / (5.0 / 3.0 - 1.0) + ], + axis=-1) + write_gkyl_field(out_dir / "shock_tube_1d_p0.gkyl", [100], [0.0], [1.0], + np.sqrt(2.0) * moments, + poly_order=0, + basis_type="serendipity") + growth_time = np.linspace(0.0, 10.0, 101) + write_gkyl_dynvector(out_dir / "exponential_energy.gkyl", growth_time, + (1e-6 * np.exp(0.4 * growth_time))[:, None]) + + # Time-resolved, positive travelling waves, encoded as p0 modal fields. + # Zero-padded frame names preserve time order under shell glob expansion. + wave_x = (np.arange(64) + 0.5) * (2 * np.pi / 64) + for frame, time in enumerate(np.linspace(0.0, 2 * np.pi, 16, endpoint=False)): + wave = 1.0 + 0.6 * np.cos(wave_x - time) + write_gkyl_field(out_dir / f"travelling_wave_{frame:03d}.gkyl", [64], [0.0], + [2 * np.pi], + np.sqrt(2) * wave[:, None], + poly_order=0, + basis_type="serendipity", + time=float(time), + frame=frame) + surface_axes = [(np.arange(32) + 0.5) * (2 * np.pi / 32)] * 2 + surface_x, surface_y = np.meshgrid(*surface_axes, indexing="ij") + for frame, time in enumerate(np.linspace(0.0, 2 * np.pi, 12, endpoint=False)): + wave = 1.0 + 0.6 * np.cos(surface_x - time) * np.cos(surface_y) + write_gkyl_field(out_dir / f"wave_surface_{frame:03d}.gkyl", [32, 32], + [0.0, 0.0], [2 * np.pi, 2 * np.pi], + 2 * wave[..., None], + poly_order=0, + basis_type="serendipity", + time=float(time), + frame=frame) + # An anisotropic Gaussian scalar field for volume and isosurface views. + volume_axis = (np.arange(24) + 0.5) / 6 - 2 + vx, vy, vz = np.meshgrid(volume_axis, volume_axis, volume_axis, indexing="ij") + blob = np.exp(-(vx**2 + 2 * vy**2 + 0.5 * vz**2)) + write_gkyl_field(out_dir / "gaussian_volume.gkyl", [24, 24, 24], [-2.0] * 3, + [2.0] * 3, + np.sqrt(8) * blob[..., None], + poly_order=0, + basis_type="serendipity") + + # --- field files (random DG coefficients) --- + for stem, ndim, cells, poly_order, basis_type in _FIELD_CONFIGS: + nc = num_comps(basis_type, ndim, poly_order) + lower = [0.0] * ndim + upper = [1.0] * ndim + values = _RNG.standard_normal((*cells, nc)) + write_gkyl_field( + out_dir / f"{stem}.gkyl", + cells, + lower, + upper, + values, + poly_order=poly_order, + basis_type=basis_type, + ) + + # --- c2p mapping files (analytical DG coordinate coefficients) --- + for entry in _C2P_CONFIGS: + stem, kind, cells, poly_order, basis_type, *extra = entry + nc_per_dim = num_comps(basis_type, len(cells), poly_order) + comp_lo = [0.0] * len(cells) + comp_hi = [1.0] * len(cells) + + if kind == "stretch": + phys_lo, phys_hi = extra + values = _c2p_stretch_values(cells, phys_lo, phys_hi, nc_per_dim) + elif kind == "rotation": + angle = extra[0] + values = _c2p_rotation_values(cells, comp_lo, comp_hi, angle, nc_per_dim) + else: + raise ValueError(f"Unknown c2p kind: {kind!r}") + + write_gkyl_field( + out_dir / f"{stem}.gkyl", + cells, + comp_lo, + comp_hi, + values, + poly_order=poly_order, + basis_type=basis_type, + ) + + # --- symmetric four-component beam profiles for a convergence plot --- + profile_lower, profile_upper, profile_cells = -2.5, 2.5, 256 + for alpha, stem in ( + (2.0e-4, "mirror_comparison_2em4_1d_ms_p1"), + (2.0e-5, "mirror_comparison_2em5_1d_ms_p1"), + ): + values = _project_1d_p1( + _mirror_comparison_profiles, + profile_lower, + profile_upper, + profile_cells, + alpha, + ) + write_gkyl_field( + out_dir / f"{stem}.gkyl", + [profile_cells], + [profile_lower], + [profile_upper], + values, + poly_order=1, + basis_type="serendipity", + ) + + # --- two-frame distribution-like family (shared grid, one file per frame) --- + distf_cells = [64, 32] + distf_nc = num_comps("serendipity", 2, 2) + for frame in (0, 1): + values = _RNG.standard_normal((*distf_cells, distf_nc)) + write_gkyl_field( + out_dir / f"distf_p2_{frame}.gkyl", + distf_cells, + [0.0, 0.0], + [1.0, 1.0], + values, + poly_order=2, + basis_type="serendipity", + time=0.1 * frame, + frame=frame, + ) + + # --- multiblock family: 3 blocks x 2 frames of one 2-D field --- + # Gkeyll's multiblock naming is '_b-_.gkyl' (a + # real example: rt_gk_multib_sheath_1x2v_p1_b2-geo_int_B3.gkyl). The + # blocks tile the x axis into abutting, disjoint domains -- one field on + # a decomposed domain, which is what postgkyl must draw as one picture. + mb_cells = [8, 6] + mb_nc = num_comps("serendipity", 2, 1) + for block in range(3): + for frame in (0, 1): + values = _RNG.standard_normal((*mb_cells, mb_nc)) + block + write_gkyl_field( + out_dir / f"mb_sim_b{block}-elc_M0_{frame}.gkyl", + mb_cells, + [float(block), 0.0], + [float(block + 1), 1.0], + values, + poly_order=1, + basis_type="serendipity", + time=0.1 * frame, + frame=frame, + ) + + # --- dynvector (bare time series, e.g. a field-energy history) --- + # Two components, each a smooth logistic growth-then-saturate curve (one + # slower-rising than the other) -- long enough (>15700 points) to + # exercise the CLI's fit/growth leading-window scan, strictly positive + # throughout (exp2's log-linear auto-guess needs y > 0), and with a + # second component so column-selecting verbs (e.g. val2coord) have + # something to select. + n_energy = 15714 + t = np.linspace(0.0, 100.0, n_energy) + y0, y_plateau, t_mid, k = 1e-6, 1.0, 30.0, 0.5 + energy0 = y0 + (y_plateau - y0) / (1.0 + np.exp(-k * (t - t_mid))) + energy1 = y0 + (y_plateau - y0) / (1.0 + np.exp(-0.5 * k * (t - 1.5 * t_mid))) + write_gkyl_dynvector(out_dir / "energy_dynvec.gkyl", t, + np.stack([energy0, energy1], axis=-1)) + + +if __name__ == "__main__": + out = Path(__file__).parent / "test_data" / "generated" + generate_all(out) + files = sorted(out.glob("*.gkyl")) + print(f"Generated {len(files)} files in {out}:") + for f in files: + print(f" {f.name}") diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py new file mode 100644 index 00000000..efa5ea83 --- /dev/null +++ b/tests/test_cli_commands.py @@ -0,0 +1,186 @@ +"""Integration tests for the uniformly generated chained CLI.""" + +from __future__ import annotations + +from pathlib import Path + +from click.testing import CliRunner +import numpy as np +import pytest + +import postgkyl as pg + +from postgkyl.cli.app import ( + COMMAND_ALIASES, + COMMANDS, + COMMAND_SECTIONS, + MODELS, + cli, +) + +ROOT = Path(__file__).parents[1] +DATA = ROOT / "tests" / "test_data" +FIELD = DATA / "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl" +DISTF = DATA / "generated" / "distf_p2_0.gkyl" +ENERGY = DATA / "generated" / "energy_dynvec.gkyl" +FIELD_3D = DATA / "generated" / "3d_ms_p1.gkyl" + + +def _run(*args): + return CliRunner().invoke(cli, [str(arg) for arg in args]) + + +def _ok(*args): + result = _run(*args) + assert result.exit_code == 0, result.output + return result + + +def test_every_registered_subcommand_is_generated(): + assert len(COMMANDS) == len(MODELS) + assert {command.name + for command in COMMANDS} == {model.name + for model in MODELS} + assert set(cli.commands) == {model.name for model in MODELS} + + +def test_help_groups_the_generated_inventory(): + result = _ok("--help") + for section, names in COMMAND_SECTIONS.items(): + assert f"{section}:" in result.output + assert names + assert "rotations_bparrotate" in result.output + assert "local_poly" in result.output + + +def test_removed_manual_commands_and_legacy_spellings_are_rejected(): + for spelling in ( + "bparrotate", + "euler", + "tenmoment", + "status", + "dg_local_poly", + "gk-load-quantity", + "extractinput", + "plotly-animate", + ): + assert _run(spelling).exit_code != 0 + + +def test_aliases_only_add_spellings(): + assert dict(COMMAND_ALIASES) == {"pl": "plot", "ev": "evaluate"} + assert cli.get_command(None, "pl") is cli.get_command(None, "plot") + assert cli.get_command(None, "ev") is cli.get_command(None, "evaluate") + + +def test_print_values_preserves_precision_and_pipeline(): + options = np.get_printoptions() + expected = np.array2string(pg.load(ENERGY).values.squeeze(), precision=16) + result = _ok(ENERGY, "print", "info") + assert result.output == expected + "\n" + _ok(ENERGY, "info").output + assert np.get_printoptions() == options + assert "print" in COMMAND_SECTIONS["Utility"] + + +@pytest.mark.parametrize("flag", ["--grid", "-g", "--grid=True"]) +def test_print_grid_axes(flag): + data = pg.load(DISTF) + expected = "".join( + np.array2string(axis, precision=16) + "\n" for axis in data.grid) + assert _ok(DISTF, "print", flag).output == expected + + +def test_print_tag_selection_and_explicit_false(): + result = _ok(ENERGY, "--tag", "energy", DISTF, "--tag", "distribution", + "print", "--use", "energy", "--grid", "False") + assert result.output == _ok(ENERGY, "print").output + assert _ok(ENERGY, "print", "--use", "missing").output == "" + assert _ok(ENERGY, ENERGY, "print").output == 2 * _ok(ENERGY, "print").output + + +def test_print_modal_coefficients_and_interpolated_values(): + for pipeline in ((), ("interpolate", )): + data = pg.load(DISTF) + if pipeline: + data = data.interpolate() + expected = np.array2string(data.values.squeeze(), precision=16) + "\n" + assert _ok(DISTF, *pipeline, "print").output == expected + + +def test_bare_filename_is_a_spelling_for_canonical_load(): + bare = _ok(FIELD, "info") + explicit = _ok("load", "--file_name", FIELD, "info") + assert bare.output == explicit.output + + +def test_fluent_chain_uses_api_command_and_option_names(): + result = _ok(FIELD, "interpolate", "select", "--z0", "0", "--comp", "0", + "info") + assert "Number of components: 1" in result.output + + +def test_select_prioritizes_the_first_option_for_each_initial(): + select = next(command for command in COMMANDS if command.name == "select") + options = {option.name: option.opts for option in select.params} + assert options == { + "comp": ["--comp", "-c"], + "z0": ["--z0", "-z"], + "z1": ["--z1"], + "z2": ["--z2"], + "z3": ["--z3"], + "z4": ["--z4"], + "z5": ["--z5"], + "inplace": ["--inplace", "-i"], + "tag": ["--tag", "-t"], + "label": ["--label", "-l"], + } + result = _ok(FIELD, "interpolate", "select", "-c", "0", "-z", "0", "-t", + "chosen", "info") + assert result.output.startswith("(chosen#0)") + + +def test_api_underscores_are_the_only_cli_spellings(): + assert _run(DISTF, "local-poly").exit_code != 0 + assert _run("load", "--file-name", FIELD, "info").exit_code != 0 + _ok(DISTF, "local_poly", "--npoints", "3", "info") + + +def test_boolean_options_accept_an_optional_explicit_value(): + _ok(DISTF, "interpolate", "fft", "--psd", "info") + _ok(DISTF, "interpolate", "fft", "--psd", "True", "info") + _ok(DISTF, "interpolate", "fft", "--psd", "False", "info") + + +def test_declared_cli_arguments_are_positional(): + _ok(ENERGY, "ev", "f0 2 *", "info") + assert _run(ENERGY, "evaluate", "--chain", "f 2 *").exit_code != 0 + + _ok(FIELD_3D, "integrate", "2", "info") + assert _run(FIELD_3D, "integrate", "--axis", "2").exit_code != 0 + assert _run(FIELD_3D, "integrate_axis", "2").exit_code != 0 + + +def test_generated_save_options_match_python_parameter_names(tmp_path): + output = tmp_path / "field" + _ok(DISTF, "save", "--out_name", output, "--extension", "npy") + assert _run(DISTF, "save", "--out-name", output).exit_code != 0 + assert output.with_suffix(".npy").is_file() + assert _run(DISTF, "save", "--out", output).exit_code != 0 + + +def test_plot_uses_generated_render_options(tmp_path): + output = tmp_path / "field.png" + _ok(FIELD, "interpolate", "select", "--comp", "0", "plot", "--no_show", + "--grid_indices", "True", "--saveas", output) + assert output.is_file() + + +def test_manual_session_render_options_are_not_registered(): + assert _run("--batch_mode", FIELD, "info").exit_code != 0 + assert _run("--saveframes-prefix", "frame", FIELD, "info").exit_code != 0 + + +def test_version_and_unknown_command_edges(): + version = _ok("--version") + assert "postgkyl" in version.output + assert _run("definitely-not-a-command").exit_code != 0 diff --git a/tests/test_cli_contract_edges.py b/tests/test_cli_contract_edges.py new file mode 100644 index 00000000..07cf36f1 --- /dev/null +++ b/tests/test_cli_contract_edges.py @@ -0,0 +1,832 @@ +"""Edge contracts for the generated CLI schema and runtime.""" + +from __future__ import annotations + +from enum import Enum +from types import ModuleType, SimpleNamespace +from typing import Annotated, Literal + +import click +from click.testing import CliRunner +import pytest + +import postgkyl.cli.compiler as compiler +from postgkyl.cli.compiler import ( + CodecKind, + CommandCompilationError, + TypeCodec, + build_click_command, + compile_callable, + compile_public_surface, + group_by_section, +) +from postgkyl.cli.docstrings import DocstringError, parse_docstring +from postgkyl.cli.state import DataSpace +from postgkyl.cli_spec import ( + ChoiceProvider, + CliArgument, + CliHidden, + CliType, + CommandSpec, + DatasetRef, + Execution, + KeyValue, + PipelineInput, + ResultPolicy, + Section, + command, + command_spec, + hidden, + hidden_spec, +) +from postgkyl.cli import app as cli_app +from postgkyl.cli.discovery import ( + SurfaceClassificationError, + _classify, + _diagnostic_modules, + discover_public_surface, +) + +pytestmark = pytest.mark.compatibility + + +class _Dataset: + """Small structural dataset used by the generic pipeline adapter.""" + + def __init__(self, tag: str): + self.tag = tag + self.ctx = {} + self.grid = [] + self.values = [] + + +class _Group: + + def __init__(self, *datasets): + self.datasets = datasets + + +class _BadEnum(Enum): + COMPOSITE = (1, 2) + + +class _Mode(Enum): + TEXT = "text" + BINARY = "binary" + + +def _compiled_option(annotation): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def option(value): + """Compile one option. + + Args: + value: Value exposed by the command. + """ + + option.__annotations__ = {"value": annotation} + return compile_callable(option) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({ + "section": "Utility", + "execution": Execution.LOAD + }, "enums"), + ({ + "section": Section.UTILITY, + "execution": "LOAD" + }, "enums"), + ({ + "section": Section.UTILITY, + "execution": Execution.LOAD, + "result": "DATA", + }, "enums"), + ({ + "section": Section.UTILITY, + "execution": Execution.COMBINE, + "consumes_inputs": 1, + }, "bool value"), + ({ + "section": Section.UTILITY, + "execution": Execution.COMBINE, + "order": True, + }, "integer"), + ({ + "section": Section.UTILITY, + "execution": Execution.MAP_REPLACE, + "consumes_inputs": True, + }, "only MAP_APPEND and COMBINE"), + ({ + "section": Section.UTILITY, + "execution": Execution.TERMINAL_ALL, + }, "non-DATA"), + ], +) +def test_command_spec_rejects_invalid_states(kwargs, message): + with pytest.raises((TypeError, ValueError), match=message): + CommandSpec(**kwargs) + + +def test_cli_marker_records_reject_invalid_values(): + with pytest.raises(ValueError, match="non-empty"): + DatasetRef("") + with pytest.raises(TypeError, match="callable"): + ChoiceProvider(1) + with pytest.raises(ValueError, match="non-empty"): + hidden(" ") + + +def test_command_and_hidden_markers_are_mutually_exclusive(): + spec = CommandSpec(Section.UTILITY, Execution.LOAD) + + def first(): + pass + + assert command(spec)(first) is first + assert command(spec)(first) is first + assert command_spec(first) == spec + assert hidden_spec(first) is None + with pytest.raises(ValueError, match="already has a CommandSpec"): + hidden("not public")(first) + with pytest.raises(ValueError, match="different CommandSpec"): + command(CommandSpec(Section.VERBS, Execution.LOAD))(first) + + @hidden("not public") + def second(): + pass + + assert hidden_spec(second).reason == "not public" + assert command_spec(second) is None + with pytest.raises(ValueError, match="already marked CliHidden"): + command(spec)(second) + with pytest.raises(TypeError, match="requires a CommandSpec"): + command("load") + + +def test_discovery_rejects_unclassified_and_doubly_classified_callables(): + + def unclassified(): + pass + + with pytest.raises(SurfaceClassificationError, match="unclassified"): + _classify(unclassified, "example.unclassified") + + def doubly_classified(): + pass + + doubly_classified.__postgkyl_command_spec__ = CommandSpec( + Section.UTILITY, Execution.LOAD) + doubly_classified.__postgkyl_cli_hidden__ = CliHidden("synthetic conflict") + with pytest.raises(SurfaceClassificationError, + match="both exposed and hidden"): + _classify(doubly_classified, "example.conflict") + + +def test_discovery_handles_cycles_facade_diagnostics_and_nonfunction_variables( +): + root = ModuleType("postgkyl.diagnostics") + child = ModuleType("postgkyl.diagnostics.synthetic") + root.__all__ = ["child", "alias"] + root.child = child + root.alias = child + child.__all__ = [] + child.VARIABLES = {"not_callable": object()} + + modules = list(_diagnostic_modules(root)) + assert modules == [root, child] + + @hidden("classified at its diagnostic home") + def diagnostic_alias(): + pass + + diagnostic_alias.__module__ = child.__name__ + + class EmptySurface: + pass + + facade = SimpleNamespace(GData=EmptySurface, + GDataGroup=EmptySurface, + __all__=["diagnostic_alias"], + diagnostic_alias=diagnostic_alias, + diagnostics=root) + assert discover_public_surface(facade) == () + + +def test_group_alias_with_missing_target_falls_through(): + group = cli_app.PgkylGroup() + context = click.Context(group) + assert group.get_command(context, "pl") is None + + +def test_group_help_ignores_missing_commands_and_empty_sections(monkeypatch): + group = cli_app.PgkylGroup() + context = click.Context(group) + formatter = context.make_formatter() + monkeypatch.setattr(cli_app, "COMMAND_SECTIONS", { + "Missing": ["not_registered"], + "Empty": [], + }) + group.format_commands(context, formatter) + assert formatter.getvalue() == "" + + +def test_group_resolve_command_with_no_arguments_delegates_to_click(): + group = cli_app.PgkylGroup() + context = click.Context(group) + with pytest.raises(IndexError): + group.resolve_command(context, []) + + +@pytest.mark.parametrize( + ("doc", "message"), + [ + (None, "missing docstring"), + ("Args:", "missing first-paragraph"), + ("Summary.\n\nArgs:\n value: one\n\nArgs:\n value: two", + "duplicate Args"), + ("Summary.\n\nArgs:\n value:", "has no description"), + ("Summary.\n\nArgs:\n value: one\n value: two", "documented twice"), + ("Summary.\n\nArgs:\nnot an entry", "malformed Args entry"), + ("Summary.\n\nArgs:\n absent: no such parameter", + "absent from signature"), + ("Summary.", "is undocumented"), + ], +) +def test_docstring_parser_reports_each_invalid_contract(doc, message): + + def target(value): + pass + + target.__doc__ = doc + with pytest.raises(DocstringError, match=message): + parse_docstring(target, required={"value"}, signature_names={"value"}) + + +def test_docstring_parser_collects_multiline_text_and_stops_at_section(): + + def target(value): + """Summary line continued + on the next line. + + Longer narrative. + + Args: + value: First part. + Second part. + + Returns: + Nothing. + """ + + parsed = parse_docstring(target, + required={"value"}, + signature_names={"value"}) + assert parsed.summary == "Summary line continued on the next line." + assert parsed.long_help.endswith("Longer narrative.") + assert parsed.parameters == {"value": "First part. Second part."} + + +@pytest.mark.parametrize( + ("annotation", "message"), + [ + (Annotated[int, object()], "unsupported Annotated marker"), + (Annotated[int, + ChoiceProvider(lambda: (1, )), + ChoiceProvider(lambda: (2, ))], "duplicate Annotated"), + (Annotated[dict[str, int], KeyValue(), + KeyValue()], "duplicate Annotated"), + (Annotated[int, CliType(int), CliType(str)], "duplicate Annotated"), + (Annotated[dict[str, int], + ChoiceProvider(lambda: (1, )), + KeyValue()], "cannot be combined"), + (Annotated[int, + ChoiceProvider(lambda: "abc")], "choice provider failed"), + (Annotated[int, ChoiceProvider(lambda: ())], "returned no choices"), + (Annotated[list[int], + ChoiceProvider(lambda: (1, ))], "requires a scalar"), + (Annotated[int, ChoiceProvider(lambda: ("one", ))], "do not match"), + (Annotated[int, ChoiceProvider(lambda: (True, ))], "do not match"), + (Annotated[int, ChoiceProvider(lambda: (1, 1))], "duplicate choices"), + (Annotated[int, KeyValue()], "requires a mapping"), + (_BadEnum, "Enum values must be CLI scalars"), + (Literal[()], "Literal choices"), + (Literal[object()], "Literal choices"), + (list[()], "list must have exactly one"), + (tuple[()], "tuple must declare"), + (dict[str, int], "mapping needs Annotated"), + (Annotated[dict[()], KeyValue()], "mapping must declare"), + (Annotated[dict[str, list[int]], KeyValue()], "must be CLI scalars"), + (int | str, "unsupported union"), + (complex, "unsupported annotation"), + ], +) +def test_compiler_rejects_lossy_option_annotations(annotation, message): + with pytest.raises(CommandCompilationError, match=message): + _compiled_option(annotation) + + +def test_choice_and_variadic_tuple_codecs_round_trip(): + calls = [] + + @command( + CommandSpec(Section.UTILITY, Execution.LOAD, result=ResultPolicy.SILENT)) + def options(*, + level: Annotated[int, ChoiceProvider(lambda: (1, 2))] = 1, + values: tuple[int, ...] = ()): + """Use provider and repeated-value codecs. + + Args: + level: Registry-provided level. + values: Repeated integer values. + """ + calls.append((level, values)) + + result = CliRunner().invoke( + build_click_command(compile_callable(options)), + ["--level", "2", "--values", "3", "--values", "4"], + obj=DataSpace()) + assert result.exit_code == 0, result.output + assert calls == [(2, [3, 4])] + + +def test_mapping_codec_rejects_malformed_and_duplicate_entries(): + model = _compiled_option(Annotated[dict[str, int] | None, KeyValue()]) + command_obj = build_click_command(model) + for arguments, message in ( + (["--value", "missing-separator"], "expected key=value"), + (["--value", "=1"], "expected key=value"), + (["--value", "a=1", "--value", "a=2"], "duplicate mapping key"), + ): + result = CliRunner().invoke(command_obj, arguments, obj=DataSpace()) + assert result.exit_code != 0 + assert message in result.output + + +def test_compile_callable_reports_signature_and_name_errors(): + + def unmarked(): + pass + + with pytest.raises(CommandCompilationError, match="has no CommandSpec"): + compile_callable(unmarked) + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def keywords(**values: int): + """Reject keyword capture. + + Args: + values: Arbitrary values. + """ + + with pytest.raises(CommandCompilationError, match=r"\*\*kwargs"): + compile_callable(keywords) + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def unresolved(value: "MissingType"): # noqa: F821 + """Reject an unresolved annotation. + + Args: + value: Unresolvable value. + """ + + with pytest.raises(CommandCompilationError, match="could not be resolved"): + compile_callable(unresolved) + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def named(): + """Reject an invalid projected name.""" + + with pytest.raises(CommandCompilationError, match="invalid command name"): + compile_callable(named, name="-named") + + +@pytest.mark.parametrize( + ("annotation", "execution", "parameter", "message"), + [ + (Annotated[object, PipelineInput(), + PipelineInput()], Execution.COMBINE, "value", + "duplicate PipelineInput"), + (object, Execution.LOAD, "*value", "not a declared pipeline"), + (Annotated[object, DatasetRef(), DatasetRef()], Execution.COMBINE, + "value", "duplicate DatasetRef"), + (Annotated[object, DatasetRef(), PipelineInput()], Execution.COMBINE, + "value", "both DatasetRef and PipelineInput"), + (Annotated[str, CliArgument(), CliArgument()], Execution.COMBINE, + "value", "duplicate CliArgument"), + (Annotated[str, CliArgument(), PipelineInput()], Execution.COMBINE, + "value", "both CliArgument and PipelineInput"), + (Annotated[str, CliArgument(), DatasetRef()], Execution.COMBINE, + "value", "both CliArgument and DatasetRef"), + (Annotated[list[str], CliArgument()], Execution.LOAD, "value", + "supports scalar values only"), + ], +) +def test_compiler_rejects_invalid_pipeline_markers(annotation, execution, + parameter, message): + namespace = {} + exec( + f"def target({parameter}):\n" + " 'Compile pipeline metadata.\\n\\n Args:\\n value: Pipeline value.'\n", + namespace, + ) + target = namespace["target"] + target.__annotations__ = {"value": annotation} + command(CommandSpec(Section.UTILITY, execution))(target) + with pytest.raises(CommandCompilationError, match=message): + compile_callable(target) + + +def test_load_rejects_an_explicit_pipeline_receiver(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def loader(value: Annotated[object, PipelineInput()]): + """Reject a receiver on a loader. + + Args: + value: Invalid pipeline receiver. + """ + + with pytest.raises(CommandCompilationError, match="LOAD commands cannot"): + compile_callable(loader) + + +def test_compile_public_surface_deduplicates_aliases_and_rejects_collisions(): + + @command(CommandSpec(Section.VERBS, Execution.LOAD, order=2)) + def first(): + """First command.""" + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def second(): + """Second command.""" + + surface = [ + SimpleNamespace(name="alias", callable=first), + SimpleNamespace(name="alias", callable=first), + SimpleNamespace(name="second", callable=second), + ] + models = compile_public_surface(surface) + assert [model.name for model in models] == ["alias", "second"] + assert group_by_section(models) == { + "Verbs": ["alias"], + "Utility": ["second"], + } + + surface[-1] = SimpleNamespace(name="alias", callable=second) + with pytest.raises(CommandCompilationError, match="command name collision"): + compile_public_surface(surface) + + +def _invoke_pipeline(fn, datasets, arguments=()): + space = DataSpace(list(datasets)) + result = CliRunner().invoke(build_click_command(compile_callable(fn)), + list(arguments), + obj=space) + assert result.exit_code == 0, result.output + return result, space + + +@pytest.mark.parametrize("execution", + [Execution.MAP_REPLACE, Execution.MAP_APPEND]) +def test_map_execution_policies_transform_each_dataset(execution): + + @command( + CommandSpec(Section.VERBS, + execution, + consumes_inputs=execution is Execution.MAP_APPEND)) + def transform(data: object): + """Transform one dataset. + + Args: + data: Current dataset. + """ + return _Dataset(data.tag + "-new") + + _, space = _invoke_pipeline(transform, [_Dataset("a"), _Dataset("b")]) + assert [dataset.tag for dataset in space.datasets] == ["a-new", "b-new"] + + +def test_map_append_can_preserve_inputs_and_flatten_groups(): + + @command(CommandSpec(Section.VERBS, Execution.MAP_APPEND)) + def append(data: object): + """Append a group for each dataset. + + Args: + data: Current dataset. + """ + return _Group(_Dataset(data.tag + "-one"), _Dataset(data.tag + "-two")) + + _, space = _invoke_pipeline(append, [_Dataset("a")]) + assert [dataset.tag for dataset in space.datasets] == ["a", "a-one", "a-two"] + + +def test_map_or_terminal_replaces_data_and_prints_scalar_values(): + + @command( + CommandSpec(Section.VERBS, + Execution.MAP_OR_TERMINAL_EACH, + result=ResultPolicy.VALUE)) + def maybe(data: object): + """Return either data or a value. + + Args: + data: Current dataset. + """ + return _Dataset("replaced") if data.tag == "data" else 7 + + result, space = _invoke_pipeline(maybe, [_Dataset("data"), _Dataset("value")]) + assert result.output == "7\n" + assert [dataset.tag for dataset in space.datasets] == ["replaced", "value"] + + +def test_terminal_each_presents_nested_values_without_mutating_data(): + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_EACH, + result=ResultPolicy.VALUE)) + def inspect_one(data: object): + """Inspect one dataset. + + Args: + data: Current dataset. + """ + return [None, data, (data.tag, )] + + original = _Dataset("kept") + result, space = _invoke_pipeline(inspect_one, [original]) + assert result.output == "kept\n" + assert space.datasets == [original] + + +def test_terminal_all_silent_result_neither_prints_nor_mutates_data(): + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) + def inspect_all(data: Annotated[list[object], PipelineInput()]): + """Inspect the entire working set. + + Args: + data: Current working set. + """ + assert [dataset.tag for dataset in data] == ["kept"] + return "not printed" + + original = _Dataset("kept") + result, space = _invoke_pipeline(inspect_all, [original]) + assert result.output == "" + assert space.datasets == [original] + + +def test_load_value_adds_dataset_results_and_presents_only_values(): + + @command( + CommandSpec(Section.UTILITY, Execution.LOAD, result=ResultPolicy.VALUE)) + def load_value(): + """Load and present a mixed result.""" + return (_Dataset("loaded"), "message") + + result, space = _invoke_pipeline(load_value, []) + assert result.output == "message\n" + assert space.datasets == [] + + +def test_combine_preserves_returned_dataset_order(): + + @command(CommandSpec(Section.VERBS, Execution.COMBINE)) + def reverse(*datasets: object): + """Reverse datasets. + + Args: + datasets: Current working set. + """ + return tuple(reversed(datasets)) + + first, second = _Dataset("first"), _Dataset("second") + _, space = _invoke_pipeline(reverse, [first, second]) + assert space.datasets == [second, first] + + +def test_dataset_reference_resolves_uniquely_and_limits_consumption(): + + @command(CommandSpec(Section.VERBS, Execution.COMBINE, consumes_inputs=True)) + def take(reference: Annotated[object, DatasetRef()] = None): + """Consume one referenced dataset. + + Args: + reference: Tagged dataset to consume. + """ + return _Dataset("result") + + kept, consumed = _Dataset("kept"), _Dataset("chosen") + _, space = _invoke_pipeline(take, [kept, consumed], ["--reference", "chosen"]) + assert [dataset.tag for dataset in space.datasets] == ["kept", "result"] + + for datasets, message in (([kept], "no dataset tagged"), + ([_Dataset("chosen"), + _Dataset("chosen")], "matches 2 datasets")): + result = CliRunner().invoke(build_click_command(compile_callable(take)), + ["--reference", "chosen"], + obj=DataSpace(datasets)) + assert result.exit_code != 0 + assert message in result.output + + +def test_pipeline_errors_keep_click_errors_and_wrap_api_errors(): + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) + def fail(*datasets: object): + """Raise an API error. + + Args: + datasets: Current working set. + """ + raise ValueError("bad input") + + result = CliRunner().invoke(build_click_command(compile_callable(fail)), [], + obj=DataSpace([_Dataset("one")])) + assert result.exit_code != 0 + assert "bad input" in result.output + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) + def fail_with_click(*datasets: object): + """Raise a Click error. + + Args: + datasets: Current working set. + """ + raise click.UsageError("click input") + + with pytest.raises(click.UsageError, match="click input"): + compiler.execute_model(SimpleNamespace(obj=DataSpace([_Dataset("one")])), + compile_callable(fail_with_click), {}) + + +def test_generated_help_skips_undocumented_internal_arguments(): + argument = compiler._DocumentedArgument(["value"], help=None) + assert argument.get_help_record(None) is None + command_obj = compiler._GeneratedCommand("internal", params=[argument]) + result = CliRunner().invoke(command_obj, ["--help"]) + assert result.exit_code == 0 + assert "Arguments:" not in result.output + + +def test_unreachable_codec_kinds_fail_loudly_and_none_stays_none(): + codec = TypeCodec(CodecKind.STRING, str) + assert compiler._convert_scalar(None, codec) is None + invalid = TypeCodec("invalid", str) + with pytest.raises(AssertionError, match="invalid"): + compiler._click_scalar(invalid) + + +def test_pipeline_receiver_requires_a_concrete_annotation(): + + @command(CommandSpec(Section.VERBS, Execution.MAP_REPLACE)) + def transform(data): + """Transform data. + + Args: + data: Current dataset. + """ + + with pytest.raises(CommandCompilationError, match="concrete annotation"): + compile_callable(transform) + + +def test_self_receiver_requires_exactly_one_direct_input(): + + @command(CommandSpec(Section.VERBS, Execution.MAP_REPLACE)) + def transform(self): + """Transform one receiver. + + Args: + self: Current dataset. + """ + + model = compile_callable(transform) + with pytest.raises(click.UsageError, match="exactly one"): + compiler._call(model, [_Dataset("one"), _Dataset("two")], {}, None) + + +def test_keyword_only_pipeline_receiver_gets_the_working_set(): + received = [] + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) + def inspect(*, data: Annotated[list[object], PipelineInput()]): + """Inspect all data. + + Args: + data: Current working set. + """ + received.append(data) + + _, space = _invoke_pipeline(inspect, [_Dataset("one")]) + assert received == [space.datasets] + + +def test_optional_dataset_reference_can_be_omitted(): + + @command(CommandSpec(Section.VERBS, Execution.COMBINE)) + def inspect(reference: Annotated[object, DatasetRef()] = None): + """Inspect an optional reference. + + Args: + reference: Optional tagged dataset. + """ + assert reference is None + return () + + _invoke_pipeline(inspect, [_Dataset("one")]) + + +@pytest.mark.parametrize( + "execution", [Execution.MAP_OR_TERMINAL_EACH, Execution.TERMINAL_EACH]) +def test_each_policies_can_silence_scalar_results(execution): + + @command(CommandSpec(Section.UTILITY, execution, result=ResultPolicy.SILENT)) + def inspect(data: object): + """Inspect one dataset silently. + + Args: + data: Current dataset. + """ + return data.tag + + result, _ = _invoke_pipeline(inspect, [_Dataset("one")]) + assert result.output == "" + + +def test_combine_can_consume_unreferenced_inputs(): + + @command(CommandSpec(Section.VERBS, Execution.COMBINE, consumes_inputs=True)) + def combine(*datasets: object): + """Replace all inputs. + + Args: + datasets: Current working set. + """ + return _Dataset("combined") + + _, space = _invoke_pipeline(combine, [_Dataset("one"), _Dataset("two")]) + assert [dataset.tag for dataset in space.datasets] == ["combined"] + + +def test_combine_and_terminal_all_apply_their_result_policies(): + + @command( + CommandSpec(Section.VERBS, Execution.COMBINE, result=ResultPolicy.VALUE)) + def combine(*datasets: object): + """Present a combined value. + + Args: + datasets: Current working set. + """ + return len(datasets) + + result, _ = _invoke_pipeline(combine, [_Dataset("one")]) + assert result.output == "1\n" + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.SILENT)) + def silent(*datasets: object): + """Ignore a terminal value. + + Args: + datasets: Current working set. + """ + return len(datasets) + + result, _ = _invoke_pipeline(silent, [_Dataset("one")]) + assert result.output == "" + + +def test_optional_enum_argument_lowers_its_default_value(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def choose(mode: Annotated[_Mode, CliArgument()] = _Mode.TEXT): + """Choose a mode. + + Args: + mode: Optional mode. + """ + + command_obj = build_click_command(compile_callable(choose)) + assert command_obj.params[0].default == "text" diff --git a/tests/test_cli_diagnostics.py b/tests/test_cli_diagnostics.py new file mode 100644 index 00000000..b0cebe79 --- /dev/null +++ b/tests/test_cli_diagnostics.py @@ -0,0 +1,202 @@ +"""Generated CLI coverage for equation-specific diagnostics.""" + +from __future__ import annotations + +import click +from click.testing import CliRunner +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl.cli.app import cli, COMMANDS +from postgkyl.cli.state import DataSpace +from postgkyl import diagnostics +from postgkyl.diagnostics.gk import distf +from postgkyl.gdata.gdata import GData + +GRID1D = [np.array([0.0, 1.0])] +COMMAND_BY_NAME = {command.name: command for command in COMMANDS} + + +def _make(values, tag="default"): + data = GData(tag=tag) + data.push(GRID1D, np.asarray(values)) + return data + + +def _invoke(name, datasets, **kwargs): + command = COMMAND_BY_NAME[name] + space = DataSpace(datasets=list(datasets)) + with click.Context(command, obj=space) as context: + context.invoke(command, **kwargs) + return space + + +def test_diagnostics_follow_gkeyll_model_families(): + assert diagnostics.__all__ == ["gk", "vm", "mom", "pkpm", "discovery"] + assert diagnostics.mom.five_moment.__name__.endswith(".mom.five_moment") + assert diagnostics.mom.enstrophy.__name__.endswith(".mom.enstrophy") + assert diagnostics.vm.kinetic.__name__.endswith(".vm.kinetic") + assert diagnostics.vm.trajectory.__name__.endswith(".vm.trajectory") + for old_name in ("gyrokinetics", "vlasov", "moments"): + assert not hasattr(diagnostics, old_name) + + +def test_gyrokinetic_diagnostics_use_concise_python_names(): + functions = ( + diagnostics.gk.energy_balance, + diagnostics.gk.nodes, + diagnostics.gk.particle_balance, + diagnostics.gk.load_distf, + diagnostics.gk.load_quantity, + ) + assert tuple(function.__name__ for function in functions) == ( + "energy_balance", + "nodes", + "particle_balance", + "load_distf", + "load_quantity", + ) + assert pg.gk is diagnostics.gk + assert pg.gk.available_quantities() == (diagnostics.gk.available_quantities()) + for bare_name in ("load_distf", "load_quantity", "available_gk_quantities"): + assert not hasattr(pg, bare_name) + for old_name in ( + "gk_energy_balance", + "gk_nodes", + "gk_particle_balance", + "load_gk_distf", + "load_gk_quantity", + ): + assert not hasattr(diagnostics.gk, old_name) + assert not hasattr(pg, old_name) + + +def test_only_canonical_diagnostic_names_are_registered(): + assert "rotations_bparrotate" in COMMAND_BY_NAME + assert "five_moment_pressure" in COMMAND_BY_NAME + assert "ten_moment_agyro" in COMMAND_BY_NAME + assert "multispecies_energetics" in COMMAND_BY_NAME + assert "kinetic_transform_frame" in COMMAND_BY_NAME + assert "pkpm_laguerre_compose" in COMMAND_BY_NAME + for name in ( + "gk_energy_balance", + "gk_nodes", + "gk_particle_balance", + "gk_load_distf", + "gk_load_quantity", + ): + assert name in COMMAND_BY_NAME + for old_name in ( + "bparrotate", + "agyro", + "energetics", + "euler", + "tenmoment", + "mhd", + "transform_frame", + "laguerre_compose", + "gyrokinetics-energy-balance", + "gyrokinetics-nodes", + "gyrokinetics-particle-balance", + "gyrokinetics-load-distf", + "gyrokinetics-load-quantity", + "gk-gk-energy-balance", + "gk-gk-nodes", + "gk-gk-particle-balance", + "gk-load-gk-distf", + "gk-load-gk-quantity", + ): + assert old_name not in COMMAND_BY_NAME + + +def test_load_distf_frame_is_text_and_cli_accepts_all_frames( + tmp_path, monkeypatch): + command = COMMAND_BY_NAME["gk_load_distf"] + frame_option = next(option for option in command.params + if option.name == "frame") + assert isinstance(frame_option.type, click.types.StringParamType) + + calls = [] + + def fake_load_distf_frame(*, frame, tag, **kwargs): + calls.append(frame) + data = GData(tag=tag, ctx={"frame": frame}) + data.push([np.array([0.0, 1.0])], np.array([[float(frame)]])) + return data + + monkeypatch.setattr(distf, "_load_distf_frame", fake_load_distf_frame) + for frame in (0, 2): + (tmp_path / f"sim-ion_fdot_{frame}.gkyl").touch() + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, [ + "gk_load_distf", + "--name", + "sim", + "--species", + "ion", + "--frame", + ":", + "--suffix", + "fdot", + ]) + assert result.exit_code == 0, result.output + assert calls == [0, 2] + + +def test_bparrotate_is_compiled_directly_from_the_script_callable(): + array = _make([[1.0, 0.0, 0.0]], tag="array") + field = _make([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], tag="field") + + space = _invoke("rotations_bparrotate", [array, field], + array="array", + field="field", + inplace=False, + tag="parallel", + label=None) + + assert len(space.datasets) == 1 + assert space.datasets[0].tag == "parallel" + np.testing.assert_allclose(space.datasets[0].values, [[1.0, 0.0, 0.0]]) + + +def test_dataset_parameters_are_tag_options_with_exact_api_names(): + command = COMMAND_BY_NAME["rotations_bparrotate"] + assert {option.opts[0] + for option in command.params} == { + "--array", + "--field", + "--inplace", + "--tag", + "--label", + } + + +def test_generated_map_diagnostic_uses_the_callable_signature(): + gamma = 5.0 / 3.0 + rho, velocity, pressure = 2.0, 0.5, 0.8 + energy = pressure / (gamma - 1.0) + 0.5 * rho * velocity**2 + moments = _make([[rho, rho * velocity, 0.0, 0.0, energy]]) + + space = _invoke("five_moment_pressure", [moments], + gas_gamma=gamma, + num_moms=None, + inplace=False, + tag="pressure", + label=None) + + assert len(space.datasets) == 1 + assert space.datasets[0].tag == "pressure" + np.testing.assert_allclose(space.datasets[0].values, [[pressure]]) + + +def test_missing_dataset_tag_fails_closed(): + array = _make([[1.0, 0.0, 0.0]], tag="array") + with pytest.raises(click.UsageError, match="field"): + _invoke("rotations_bparrotate", [array], + array="array", + field="field", + inplace=False, + tag=None, + label=None) diff --git a/tests/test_cli_generator.py b/tests/test_cli_generator.py new file mode 100644 index 00000000..cd0fd81a --- /dev/null +++ b/tests/test_cli_generator.py @@ -0,0 +1,446 @@ +"""Schema, compilation, discovery, and generated invocation contracts.""" + +from __future__ import annotations + +import ast +from enum import Enum +from pathlib import Path +from typing import Annotated, Any, Literal + +import click +from click.testing import CliRunner +import pytest + +import postgkyl.cli.compiler as compiler +from postgkyl.cli_spec import ( + CliArgument, + CliType, + CommandSpec, + Execution, + KeyValue, + PipelineInput, + ResultPolicy, + Section, + command, +) +from postgkyl.cli.app import COMMANDS, MODELS +from postgkyl.cli.compiler import ( + CodecKind, + CommandCompilationError, + build_click_command, + compile_callable, +) +from postgkyl.cli.discovery import discover_public_surface +from postgkyl.cli.docstrings import DocstringError +from postgkyl.cli.state import DataSpace +from postgkyl.operations import average + + +class Format(Enum): + TEXT = "text" + BINARY = "binary" + + +_CALLS = [] + + +@command( + CommandSpec(Section.UTILITY, Execution.LOAD, result=ResultPolicy.SILENT)) +def codec_demo(required: int, + *, + optional: str | None = None, + enabled: bool = False, + mode: Literal["a", "b"] = "a", + format: Format = Format.TEXT, + paths: list[Path] = [], + pair: tuple[int, float] = (1, 2.0), + values: Annotated[dict[str, int] | None, + KeyValue()] = None): + """Exercise every lossless command codec. + + Args: + required: Required integer value. + optional: Optional string value. + enabled: Explicit boolean value. + mode: Literal mode. + format: Output format. + paths: Repeatable filesystem paths. + pair: Fixed integer/float pair. + values: Repeatable key/value mapping. + """ + call = (required, optional, enabled, mode, format, paths, pair, values) + _CALLS.append(call) + return call + + +def test_codec_models_and_round_trip(tmp_path): + model = compile_callable(codec_demo) + assert model.name == "codec_demo" + by_name = {parameter.name: parameter for parameter in model.parameters} + assert by_name["required"].required + assert by_name["enabled"].codec.kind is CodecKind.BOOLEAN + assert by_name["mode"].codec.choices == ("a", "b") + assert by_name["paths"].codec.multiple + assert by_name["pair"].codec.nargs == 2 + assert by_name["values"].codec.kind is CodecKind.MAPPING + + path = tmp_path / "x" + result = CliRunner().invoke(build_click_command(model), [ + "--required", + "4", + "--optional", + "x", + "--enabled", + "True", + "--mode", + "b", + "--format", + "binary", + "--paths", + str(path), + "--pair", + "2", + "3.5", + "--values", + "n=7", + ], + obj=DataSpace()) + assert result.exit_code == 0, result.output + assert _CALLS[-1] == (4, "x", True, "b", Format.BINARY, [path], (2, 3.5), { + "n": 7 + }) + + +def test_boolean_options_are_optional_value_flags_with_false_defaults(): + calls = [] + + @command( + CommandSpec(Section.UTILITY, Execution.LOAD, result=ResultPolicy.SILENT)) + def booleans(*, enabled: bool = False): + """Exercise an optional-value boolean option. + + Args: + enabled: Enable the optional behavior. + """ + calls.append(enabled) + + model = compile_callable(booleans) + command_obj = build_click_command(model) + options = {option.name: option for option in command_obj.params} + assert options["enabled"].opts == ["--enabled", "-e"] + assert options["enabled"].default is False + + runner = CliRunner() + for arguments, expected in ( + ([], False), + (["--enabled"], True), + (["--enabled", "True"], True), + (["--enabled", "False"], False), + (["-e"], True), + ): + result = runner.invoke(command_obj, arguments, obj=DataSpace()) + assert result.exit_code == 0, result.output + assert calls[-1] == expected + + +def test_boolean_options_must_default_to_false(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def default_true(enabled: bool = True): + """Declare an invalid true-default boolean. + + Args: + enabled: Invalid true-default boolean. + """ + + with pytest.raises(CommandCompilationError, + match="boolean CLI options must default to False"): + compile_callable(default_true) + + +def test_exact_option_projection_and_help_provenance(): + model = compile_callable(codec_demo) + command_obj = build_click_command(model) + assert {option.opts[0] + for option in command_obj.params} == { + "--required", + "--optional", + "--enabled", + "--mode", + "--format", + "--paths", + "--pair", + "--values", + } + options = {option.name: option.opts for option in command_obj.params} + assert options == { + "required": ["--required", "-r"], + "optional": ["--optional", "-o"], + "enabled": ["--enabled", "-e"], + "mode": ["--mode", "-m"], + "format": ["--format", "-f"], + "paths": ["--paths", "-p"], + "pair": ["--pair"], + "values": ["--values", "-v"], + } + assert next(p for p in command_obj.params if p.name == "required").help \ + == "Required integer value." + assert command_obj.short_help == "Exercise every lossless command codec." + + +def test_short_options_prioritize_parameter_order_and_reserve_help(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def collisions(*, + alpha: int = 0, + another: int = 0, + beta_value: int = 0, + hidden: int = 0): + """Exercise shorthand conflict handling. + + Args: + alpha: First conflicting option. + another: Second conflicting option. + beta_value: Unambiguous option. + hidden: Option whose initial belongs to help. + """ + + command_obj = build_click_command(compile_callable(collisions)) + options = {option.name: option.opts for option in command_obj.params} + assert options == { + "alpha": ["--alpha", "-a"], + "another": ["--another"], + "beta_value": ["--beta_value", "-b"], + "hidden": ["--hidden"], + } + + +def test_cli_type_projects_a_broader_python_option(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def projected(*, + output: Annotated[str | list[str] | None, + CliType(str | None)] = None): + """Project direct-Python options explicitly. + + Args: + output: One command-line output path. + """ + + model = compile_callable(projected) + assert [parameter.name for parameter in model.parameters] == ["output"] + assert model.parameters[0].codec.python_type is str + + +def test_strict_compilation_rejects_missing_docs_and_unknown_types(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def undocumented(value: int): + """Missing parameter documentation.""" + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def any_value(value: Any): + """An unsupported value. + + Args: + value: Arbitrary value. + """ + + with pytest.raises(DocstringError, match="value"): + compile_callable(undocumented) + with pytest.raises(CommandCompilationError, match="Any"): + compile_callable(any_value) + + +def test_optional_annotation_does_not_make_a_required_option_optional(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def required_optional(value: str | None): + """Accept an explicitly nullable required value. + + Args: + value: Required value which may be null in direct Python calls. + """ + + model = compile_callable(required_optional) + assert model.parameters[0].required + + +def test_cli_argument_marker_projects_only_the_declared_parameter_positionally( +): + calls = [] + + @command( + CommandSpec(Section.UTILITY, Execution.LOAD, result=ResultPolicy.SILENT)) + def positional(value: Annotated[str, CliArgument()], *, suffix: str = ""): + """Accept one positional CLI value. + + Args: + value: Required positional value. + suffix: Optional suffix. + """ + calls.append(value + suffix) + + command_obj = build_click_command(compile_callable(positional)) + assert isinstance(command_obj.params[0], click.Argument) + assert isinstance(command_obj.params[1], click.Option) + help_text = CliRunner().invoke(command_obj, ["--help"]).output + assert "Arguments:" in help_text + assert "VALUE Required positional value." in help_text + result = CliRunner().invoke(command_obj, ["hello", "--suffix", "!"], + obj=DataSpace()) + assert result.exit_code == 0, result.output + assert calls == ["hello!"] + assert CliRunner().invoke(command_obj, ["--value", "hello"], + obj=DataSpace()).exit_code != 0 + + +def test_cli_argument_marker_requires_a_positional_python_parameter(): + + @command(CommandSpec(Section.UTILITY, Execution.LOAD)) + def invalid(*, value: Annotated[str, CliArgument()]): + """Reject a positional CLI marker on a keyword-only API parameter. + + Args: + value: Invalid positional projection. + """ + + with pytest.raises(CommandCompilationError, match="positional Python"): + compile_callable(invalid) + + +def test_concrete_annotated_alias_is_not_reprocessed(monkeypatch): + """Keep runtime CLI metadata authoritative over resolved string hints.""" + original = compiler.get_type_hints + evaluated = None + + def track_evaluated(source, **kwargs): + nonlocal evaluated + evaluated = source.__annotations__ + return original(source, **kwargs) + + monkeypatch.setattr(compiler, "get_type_hints", track_evaluated) + + model = compile_callable(average) + weight = next(parameter for parameter in model.parameters + if parameter.name == "weight") + assert "weight" not in evaluated + assert weight.dataset_ref + assert weight.codec.optional + + +def test_public_inventory_is_total_unique_and_deterministic(): + first = discover_public_surface() + second = discover_public_surface() + assert first == second + assert len({model.name for model in MODELS}) == len(MODELS) + assert {command_obj.name + for command_obj in COMMANDS} >= { + "interpolate", + "five_moment_pressure", + "plot", + "load", + } + assert {command_obj.name + for command_obj in COMMANDS} == {model.name + for model in MODELS} + + +def test_command_spec_rejects_invalid_loader_state(): + with pytest.raises(ValueError, match="LOAD"): + CommandSpec(Section.UTILITY, Execution.LOAD, consumes_inputs=True) + + +def test_pipeline_input_adapter_receives_the_working_set(): + calls = [] + + @command( + CommandSpec(Section.UTILITY, + Execution.TERMINAL_ALL, + result=ResultPolicy.VALUE)) + def terminal(data: Annotated[list[object], PipelineInput()], + *, + value: int = 1): + """Inspect the pipeline input. + + Args: + data: Injected command-line working set. + value: Value to return. + """ + calls.append(data) + return value + + member = object() + space = DataSpace(datasets=[member]) + result = CliRunner().invoke(build_click_command(compile_callable(terminal)), + ["--value", "7"], + obj=space) + assert result.exit_code == 0, result.output + assert calls == [[member]] + assert result.output == "7\n" + + +def test_no_scientific_click_decorators_remain(): + allowed = {"app.py"} + root = Path(__file__).parents[1] / "src" / "postgkyl" / "cli" + offenders = [] + for path in root.rglob("*.py"): + tree = ast.parse(path.read_text(), path) + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance( + node.func, ast.Attribute): + continue + if isinstance(node.func.value, ast.Name) and node.func.value.id == "click" \ + and node.func.attr in {"command", "option", "argument"} \ + and path.name not in allowed: + offenders.append(str(path.relative_to(root))) + assert offenders == [] + + +def test_every_generated_name_preserves_api_underscores(): + discovered = discover_public_surface() + assert {item.name for item in discovered} == {model.name for model in MODELS} + for item in discovered: + assert "-" not in item.name + for model, command_obj in zip(MODELS, COMMANDS): + assert command_obj.name == model.name + options = {parameter.name: parameter for parameter in command_obj.params} + expected = { + parameter.name + for parameter in model.parameters if not parameter.injected + } + assert set(options) == expected + claimed_initials = {"h"} + for parameter in model.parameters: + if not parameter.injected: + if parameter.argument: + assert isinstance(options[parameter.name], click.Argument) + expected_opts = [parameter.name] + else: + assert isinstance(options[parameter.name], click.Option) + expected_opts = ["--" + parameter.name] + initial = parameter.name[0] + if initial not in claimed_initials: + expected_opts.append("-" + initial) + claimed_initials.add(initial) + assert options[parameter.name].opts == expected_opts + + +def test_every_generated_boolean_has_a_false_cli_default(): + for model, command_obj in zip(MODELS, COMMANDS): + options = {parameter.name: parameter for parameter in command_obj.params} + for parameter in model.parameters: + if parameter.codec is None or parameter.codec.kind is not CodecKind.BOOLEAN: + continue + option = options[parameter.name] + assert option.default is False + assert option.metavar == "[BOOLEAN]" + assert parameter.default is False + assert option.opts[0] == "--" + parameter.name + + +def test_cli_has_no_manual_command_package_or_compatibility_layer(): + root = Path(__file__).parents[1] / "src" / "postgkyl" / "cli" + assert not (root / "commands").exists() + assert not (root / "compat.py").exists() + assert not (root / "legacy.py").exists() diff --git a/tests/test_commands.py b/tests/test_commands.py deleted file mode 100644 index 1fb8251c..00000000 --- a/tests/test_commands.py +++ /dev/null @@ -1,229 +0,0 @@ -"""Postgkyl module for testing click commands.""" -import click -import importlib.util -import matplotlib.pyplot as plt -import numpy as np -import os -import pytest -import subprocess - -import postgkyl.commands as cmd -from postgkyl.pgkyl import cli - -class TestCommands: - """Base class for testing Postgkyl commands. - - Note that commands which just wrap other Postgkyl functions are not tested thoroughly, - the goal here is to test if the command runs at all; more thorough testing should be - delegated to the functions themselves. - """ - dir_path = f"{os.path.dirname(__file__)}/test_data" - - ctx = click.core.Context(cli) - ctx.obj = {} - ctx.obj["in_data_strings"] = [f"{dir_path:s}/twostream-f-p2.gkyl", f"{dir_path:s}/twostream-f-p2.gkyl", f"{dir_path:s}/twostream-f-p2_0.bp"] - ctx.obj["in_data_strings_loaded"] = 0 - ctx.obj["verbose"] = False - ctx.obj["data"] = cmd.DataSpace() - - ctx.obj["fig"] = "" - ctx.obj["ax"] = "" - - ctx.obj["compgrid"] = None - ctx.obj["global_var_names"] = None - ctx.obj["global_cuts"] = (None, None, None, None, None, None, None) - ctx.obj["global_c2p"] = None - ctx.obj["global_c2p_vel"] = None - - ctx.obj["rcParams"] = {} - - # Check if ADIOS is isntalled - adios_loader = importlib.util.find_spec('adios2') - adios_missing = adios_loader is None - - # Check if ffmpeg is installed - ffmpeg_missing = True - try: - subprocess.run("ffmpeg") - ffmpeg_missing = False - except FileNotFoundError: - ffmpeg_missing = True - # end - - def test_load(self): - self.ctx.invoke(cmd.load) - data = self.ctx.obj['data'].get_dataset(0) - num_cells = data.num_cells - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_array_equal(num_cells, (64, 32)) - - - def test_ev_gkyl(self): - # Check baseline addition - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.ev, chain='f[0] f[0] +') - data = self.ctx.obj['data'].get_dataset(0) - values = data.get_values() - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_approx_equal(np.max(values), 3.352029) - - # Check longer chain, substraction, and not using dataset id - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.ev, chain='f f + f -') - data = self.ctx.obj['data'].get_dataset(0) - values = data.get_values() - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_approx_equal(np.max(values), 1.676014) - - # Check tags - self.ctx.invoke(cmd.load, tag='ts0') - self.ctx.invoke(cmd.load, tag='ts1') - self.ctx.invoke(cmd.ev, chain='ts0 ts0 +') - data = self.ctx.obj['data'].get_dataset(0, tag='ts0') - values = data.get_values() - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_approx_equal(np.max(values), 3.3520293) - - # Check ev functionality on multiple dataset together - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.ev, chain='f[:] 2 *') - data0 = self.ctx.obj['data'].get_dataset(0) - values0 = data0.get_values() - data1 = self.ctx.obj['data'].get_dataset(1) - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - values1 = data1.get_values() - np.testing.assert_approx_equal(np.max(values0), 3.3520293) - np.testing.assert_approx_equal(np.max(values1), 3.3520293) - - - @pytest.mark.skipif(adios_missing, reason="ADIOS2 is not installed") - def test_ev_adios(self): - # Check metadata - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.ev, chain='f[2] f[2].charge *') - data = self.ctx.obj['data'].get_dataset(2) - values = data.get_values() - charge = data.ctx["charge"] - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_approx_equal(np.min(values), -1.676014) - # Check if metadata is properly passed through ev: - np.testing.assert_approx_equal(charge, -1.0) - - - def test_interpolate(self): - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.interpolate) - data = self.ctx.obj['data'].get_dataset(0) - num_cells = data.num_cells - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_array_equal(num_cells, (192, 96)) - - - def test_select(self): - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.select, z0='0:10', z1='0.0', comp='0,3') - data = self.ctx.obj['data'].get_dataset(0) - values_shape = data.values.shape - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_array_equal(values_shape, (10, 1, 2)) - - - def test_plot(self): - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.plot, show=False) - fig = plt.gcf() - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - label = fig.figure.get_supylabel() - plt.close("all") - assert label == "$z_1$" - - def test_animate_save_gif(self, tmp_path): - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.load) - fn = tmp_path / "test_anim.gif" - self.ctx.invoke(cmd.animate, show=False, saveas=fn) - fig = plt.gcf() - label = fig.figure.get_supylabel() - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - plt.close("all") - assert label == "$z_1$" - assert fn.exists() - - @pytest.mark.skipif(ffmpeg_missing, reason="ffmpeg is not installed") - def test_animate_save_mp4(self, tmp_path): - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.load) - fn = tmp_path / "test_anim.mp4" - self.ctx.invoke(cmd.animate, show=False, saveas=fn) - fig = plt.gcf() - label = fig.figure.get_supylabel() - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - plt.close("all") - assert label == "$z_1$" - assert fn.exists() - - def test_grid(self): - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.grid) - data = self.ctx.obj['data'].get_dataset(0) - values_shape = data.values.shape - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings_loaded"] = 0 - np.testing.assert_array_equal(values_shape, (65, 33, 2)) - np.testing.assert_approx_equal(np.max(data.values[...,0]), 6.283185) - np.testing.assert_approx_equal(np.max(data.values[...,1]), 6) - - - def test_gk_rz(self): - # gk-rz operates on data already loaded onto the stack and locates the - # geometry from the loaded file's prefix (here '-geo_int_mapc2p.gkyl'). - saved_strings = self.ctx.obj["in_data_strings"] - self.ctx.obj["in_data_strings"] = [f"{self.dir_path:s}/gk_ltx_iwl_2x2v_p1-elc_M2par_10.gkyl"] - self.ctx.obj["in_data_strings_loaded"] = 0 - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.gk_rz) - data = self.ctx.obj['data'].get_dataset(0, tag='rz') - values = data.values - grid = data.grid - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings"] = saved_strings - self.ctx.obj["in_data_strings_loaded"] = 0 - # The 2D field is mapped onto the R-Z plane: R and Z grids plus a single - # component dimension. - assert values is not None - assert values.shape[-1] == 1 - assert len(grid) == 2 - - def test_gk_fluxsurf(self): - # gk-fluxsurf operates on data already loaded onto the stack and locates the - # geometry from the loaded file's prefix (here '-geo_int_mapc2p.gkyl'). - saved_strings = self.ctx.obj["in_data_strings"] - self.ctx.obj["in_data_strings"] = [f"{self.dir_path:s}/rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl"] - self.ctx.obj["in_data_strings_loaded"] = 0 - self.ctx.invoke(cmd.load) - self.ctx.invoke(cmd.gk_fluxsurf, x_idx=0, nphi=4, nz_interp=1) - data = self.ctx.obj['data'].get_dataset(0, tag='fluxsurf') - values = data.values - grid = data.grid - self.ctx.obj['data'].clean() - self.ctx.obj["in_data_strings"] = saved_strings - self.ctx.obj["in_data_strings_loaded"] = 0 - # The 2D field is mapped onto the theta-phi plane: theta and phi grids plus a single - # component dimension. - assert values is not None - assert values.shape[-1] == 1 - assert len(grid) == 2 \ No newline at end of file diff --git a/tests/test_coverage_container.py b/tests/test_coverage_container.py new file mode 100644 index 00000000..34050822 --- /dev/null +++ b/tests/test_coverage_container.py @@ -0,0 +1,386 @@ +"""Coverage-completing tests for gdata/gdata, gdatastate/state, gdatastate/collection, cli/*. + +These target branches the golden-path tests in test_postgkyl.py don't reach: +state readers on empty/bare containers, the modal .mul()/.div() aliases, the +CLI's abbreviation/ambiguity/fail paths, and the write/plot command edges. + +Run: PYTHONPATH=src pytest tests/test_coverage_container.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg # noqa: E402 +from postgkyl import gpython # noqa: E402 +from postgkyl.gdatastate.gdatastate import GDataState # noqa: E402 +from postgkyl.gdatastate.collection import flatten_datasets # noqa: E402 + +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +# --------------------------------------------------------------------- gdata +@needs_gkeyll +def test_mul_div_explicit_aliases_match_operators(): + a, b = pg.load(F1), pg.load(F1) + np.testing.assert_allclose(a.mul(b).values, (a * b).values) + a2, b2 = pg.load(F1), pg.load(F1) + np.testing.assert_allclose(a2.div(b2).values, (a2 / b2).values) + + +# --------------------------------------------------------------------- tags +def test_tag_setter_ignores_falsy_value(): + d = GDataState() + d.tag = "custom" + assert d.tag == "custom" + d.tag = "" # falsy: must not clobber the existing tag + assert d.tag == "custom" + + +def test_label_getter_setter(): + d = GDataState() + assert d.label == "" + d.label = "raw-label" + assert d.label == "raw-label" + + +# ---------------------------------------------------------------- shape info +def test_num_cells_falls_back_to_values_shape_then_empty(): + d = GDataState() + assert d.num_cells.size == 0 # no ctx, no values + + d.push([np.linspace(0.0, 1.0, 4)], np.zeros((3, 2))) + del d.ctx["cells"] + assert np.array_equal(d.num_cells, [3]) + + +def test_num_comps_falls_back_to_values_shape_then_zero(): + d = GDataState() + assert d.num_comps == 0 # no ctx, no values + + d.push([np.linspace(0.0, 1.0, 4)], np.zeros((3, 2))) + del d.ctx["num_comps"] + assert d.num_comps == 2 + + +@needs_gkeyll +def test_num_comps_falls_back_for_gkyl_backed_values(): + d = pg.load(F1) + del d.ctx["num_comps"] + assert d.num_comps == d.native.ncomp + + +def test_num_dims_falls_back_to_values_ndim_then_zero(): + d = GDataState() + assert d.num_dims == 0 # no ctx, no values + + d.push([np.linspace(0.0, 1.0, 4)], np.zeros((3, 2))) + del d.ctx["cells"] + assert d.num_dims == 1 + + +def test_bounds_falls_back_to_grid_then_none(): + d = GDataState() + assert d.bounds == (None, None) + + d.push([np.linspace(0.0, 2.0, 4)], np.zeros((3, 2))) + del d.ctx["lower"] + del d.ctx["upper"] + lo, up = d.bounds + np.testing.assert_allclose(lo, [0.0]) + np.testing.assert_allclose(up, [2.0]) + + +# ------------------------------------------------ getitem / setitem / copy +def test_getitem_raises_when_empty(): + d = GDataState() + with pytest.raises(ValueError): + d[0] + + +def test_getitem_uses_numpy_axis_order_when_loaded(): + d = GDataState() + d.push([np.linspace(0.0, 1.0, 4)], np.arange(6, dtype=float).reshape(3, 2)) + np.testing.assert_allclose(d[:, 1], [1.0, 3.0, 5.0]) + np.testing.assert_allclose(d[1], [2.0, 3.0]) + + +def test_setitem_uses_numpy_axis_order_when_loaded(): + d = GDataState() + d.push([np.linspace(0.0, 1.0, 4)], np.arange(12, dtype=float).reshape(3, 4)) + d[:, 2:4] *= 5 + np.testing.assert_allclose(d.values, [ + [0.0, 1.0, 10.0, 15.0], + [4.0, 5.0, 30.0, 35.0], + [8.0, 9.0, 50.0, 55.0], + ]) + + +def test_setitem_raises_when_empty(): + d = GDataState() + with pytest.raises(ValueError, match="cannot assign"): + d[0] = 1.0 + + +@needs_gkeyll +def test_setitem_rejects_native_storage(): + d = pg.load(F1) + with pytest.raises(ValueError, match="native Gkeyll storage"): + d[0] = 1.0 + + +def test_copy_with_data_deep_copies_numpy_backend(): + d = GDataState() + d.push([np.linspace(0.0, 1.0, 4)], np.ones((3, 2))) + c = d.clone() + c.values[0, 0] = 99.0 + assert d.values[0, 0] == 1.0 + assert c.grid[0] is not d.grid[0] + + +@needs_gkeyll +def test_copy_with_data_deep_copies_gkyl_backend(): + d = pg.load(F1) + c = d.clone() + assert c.native is not d.native + np.testing.assert_allclose(c.values, d.values) + + +@needs_gkeyll +def test_result_applies_explicit_tag_and_label(): + d = pg.load(F1).interpolate(tag="custom-tag", label="custom-label") + assert d.tag == "custom-tag" + assert d.label == "custom-label" + + +def test_require_operable_raises_on_empty_dataset(): + d = GDataState() + with pytest.raises(ValueError): + d._require_operable() + + +# -------------------------------------------------------------------- info +@needs_gkeyll +def test_info_reports_nodal_and_quad_representation(): + a = pg.load(F1) + assert "nodal" in a.to_nodal().info() + assert "quad" in a.to_quad().info() + + +def test_output_identity_handles_an_unrecognized_name(monkeypatch): + monkeypatch.setattr("postgkyl.gdatastate.gdatastate.io.parse_output_name", + lambda _path: None) + d = GDataState() + d._file_name = "unrecognized" + d._stamp_output_name() + assert d.output_name is None + assert "sim" not in d.ctx + + +def test_info_reports_all_optional_metadata(capsys): + d = GDataState( + ctx={ + "time": 1.5, + "frame": 3, + "block": 2, + "sim": "demo", + "basis_type": "serendipity", + "poly_order": None, + "value_form": "quad", + "num_quad": 3, + "changeset": "abc123", + "builddate": "today", + "geometry_type": "tokamak", + "geqdsk_sign_convention": -1, + "mass": 2.0, + "charge": -1.0, + "gas_gamma": 5.0 / 3.0, + "vdim": 2, + "custom_metadata": "kept", + }) + d.push([np.linspace(0.0, 1.0, 3)], np.arange(4, dtype=float).reshape(2, 2)) + + out = d.info(no_header=True) + + assert "GEQDSK sign convention: -1" in out + assert "Adiabatic index" in out + assert "custom_metadata: kept" in out + assert "default#0" not in out + assert capsys.readouterr().out == out + "\n" + + +@pytest.mark.parametrize("ctx", [{ + "builddate": "today" +}, { + "changeset": "abc123" +}, { + "geqdsk_sign_convention": 1 +}, { + "mass": 1.0 +}, { + "charge": -1.0 +}]) +def test_info_reports_independent_optional_metadata(ctx): + d = GDataState(ctx=ctx) + d.push([np.linspace(0.0, 1.0, 3)], np.ones((2, 1))) + assert d.info() + + +def test_info_handles_values_without_a_grid_and_a_grid_without_values(): + values_only = GDataState() + values_only.values = np.ones((2, 1)) + assert "Maximum" in values_only.info(no_header=True) + assert "Grid:" not in values_only.info(no_header=True) + + grid_only = GDataState() + grid_only.grid = [np.linspace(0.0, 1.0, 3)] + assert "Grid:" in grid_only.info(no_header=True) + assert "Maximum" not in grid_only.info(no_header=True) + + +# ---------------------------------------------------------- repr/str/summary +def test_repr_and_str_on_empty_dataset(): + d = GDataState() + assert "empty" in repr(d) + assert repr(d) == str(d) + + +def test_repr_and_str_on_loaded_modal_dataset(): + d = pg.load(F1) + r = repr(d) + assert "comp" in r and "tag" in r + s = str(d) + assert s.startswith(r) + assert "modal" in r or "gkyl-native" in r + + +def test_repr_on_interpolated_dataset(): + d = pg.load(F1).interpolate() + r = repr(d) + assert "interpolate" in r + + +def test_repr_handles_values_without_grid_or_basis_metadata(): + d = GDataState() + d.values = np.ones((2, 1)) + assert "[" not in repr(d) + + +def test_repr_handles_basis_without_poly_order(): + d = GDataState(ctx={"basis_type": "serendipity"}) + d.values = np.ones((2, 1)) + assert "serendipity" in repr(d) + assert " p" not in repr(d) + + +@needs_gkeyll +def test_repr_on_nodal_and_quad_datasets(): + a = pg.load(F1) + assert "nodal" in repr(a.to_nodal()) + assert "quad" in repr(a.to_quad()) + + +# ------------------------------------------------------------- collections +def test_flatten_datasets_passes_through_non_dataset_items(): + out = flatten_datasets([1, [2, 3], "x"]) + assert out == [1, 2, 3, "x"] + + +# ------------------------------------------------------------------ cli app +def test_cli_hidden_alias_pl_resolves_to_plot(tmp_path): + from click.testing import CliRunner + from postgkyl.cli.app import cli + + out = tmp_path / "alias.png" + result = CliRunner().invoke(cli, [ + F1, "interp", "sel", "--comp", "0", "pl", "--no_show", "--saveas", + str(out) + ]) + assert result.exit_code == 0, result.output + assert out.exists() + + +def test_cli_ambiguous_abbreviation_fails(): + from click.testing import CliRunner + from postgkyl.cli.app import cli + + result = CliRunner().invoke( + cli, [F1, "in"]) # "in" prefixes both info and interpolate + assert result.exit_code != 0 + assert "Ambiguous command" in result.output + + +def test_cli_unknown_token_is_neither_command_nor_file(): + from click.testing import CliRunner + from postgkyl.cli.app import cli + + result = CliRunner().invoke(cli, ["not-a-command-or-file-xyz"]) + assert result.exit_code != 0 + assert "No such command 'not-a-command-or-file-xyz'" in result.output + + +def test_cli_plot_without_datasets_raises_usage_error(): + from click.testing import CliRunner + from postgkyl.cli.app import cli + + result = CliRunner().invoke(cli, ["plot"]) + assert result.exit_code != 0 + assert "no datasets selected" in result.output + + +def test_cli_has_no_manual_batch_mode(tmp_path, monkeypatch): + from click.testing import CliRunner + from postgkyl.cli.app import cli + + monkeypatch.chdir(tmp_path) + runner = CliRunner() + result = runner.invoke(cli, ["--batch_mode", F1, "info"]) + assert result.exit_code != 0 + assert "No such option '--batch_mode'" in result.output + + +def test_cli_module_entry_point_runs_as_script(monkeypatch, capsys): + """Exercise ``if __name__ == "__main__": cli()`` in-process (so it's + visible to coverage), rather than via subprocess.""" + import runpy + monkeypatch.setattr(sys, "argv", ["pgkyl", "--help"]) + app_path = os.path.join(SRC, "postgkyl", "cli", "app.py") + with pytest.raises(SystemExit) as exc: + runpy.run_path(app_path, run_name="__main__") + assert exc.value.code == 0 + assert "Postprocessing and plotting tool" in capsys.readouterr().out + + +def test_dataspace_is_iterable(): + from postgkyl.cli.state import DataSpace + ds = DataSpace(datasets=[1, 2, 3]) + assert list(ds) == [1, 2, 3] + + +def test_cli_save_command(tmp_path): + from click.testing import CliRunner + from postgkyl.cli.app import cli + + out = tmp_path / "written.txt" + result = CliRunner().invoke(cli, [ + F1, "interp", "sel", "--comp", "0", "save", "--out_name", + str(out), "--extension", "txt" + ]) + assert result.exit_code == 0, result.output + assert out.exists() + assert str(out) in result.output diff --git a/tests/test_coverage_io.py b/tests/test_coverage_io.py new file mode 100644 index 00000000..fa5210e0 --- /dev/null +++ b/tests/test_coverage_io.py @@ -0,0 +1,523 @@ +"""Coverage-completing tests for the ``io`` leaf layer. + +Golden-path loads in test_postgkyl.py / test_gpython_rio.py only exercise the +happy path of each reader (full, non-partial, version-1, real_type f8 field +reads). This file targets the edges: partial loads (``axes=``/``comp=``), +dynvector multi-chunk continuation, legacy version-0 / float32 files, ghost +cells, the reader-registry failure path, and every ``write()`` format. + +Run: PYTHONPATH=src pytest tests/test_coverage_io.py -v +""" + +import os +import sys +from pathlib import Path + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg # noqa: E402 +from postgkyl import gpython, io # noqa: E402 +from postgkyl.io import mapping, writer # noqa: E402 +from postgkyl.io.gkyl_reader import GkylReader # noqa: E402 +from postgkyl.io.gkyl_c_reader import GkylCReader # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F1D_SINGLE_RANGE = os.path.join(DATA, "generated", + "1d_ms_p1.gkyl") # file_type 1 +F2D = os.path.join(DATA, "generated", "2d_ms_p1.gkyl") +DYNVEC = os.path.join(DATA, "generated", "energy_dynvec.gkyl") + +# ndim=1, 24 cells, 6 comps, split across 4 multi-ranges of 6 cells each +# (1-indexed [1,6] [7,12] [13,18] [19,24]) -- verified by direct header +# inspection; used below to force one whole range to be excluded on a +# partial (``axes=``) read. + + +# --------------------------------------------------------------------- io/__init__ +def test_read_defaults_ctx_to_a_fresh_dict(): + grid, values = io.read(F1) # ctx omitted entirely + assert values is not None + + +def test_read_raises_when_no_reader_is_compatible(tmp_path): + bogus = tmp_path / "not_a_gkyl_file.dat" + bogus.write_bytes(b"nope, not a gkyl file") + with pytest.raises(NameError, match="cannot be read"): + io.read(str(bogus)) + + +# --------------------------------------------------------------- gkyl_c_reader +@needs_gkeyll +def test_gkyl_c_reader_is_compatible_swallows_backend_errors(monkeypatch): + + def _raise(*a, **k): + raise RuntimeError("simulated backend failure") + + monkeypatch.setattr(gpython.rio, "file_type", _raise) + r = GkylCReader(F1, ctx={}) + assert r.is_compatible() is False + + +@needs_gkeyll +def test_gkyl_c_reader_declines_a_partial_load_request(): + r = GkylCReader(F1, ctx={}, axes=("0", None, None, None, None, None)) + assert r.is_compatible() is False + + +@needs_gkeyll +def test_gkyl_c_reader_rejects_cell_array_mismatch(monkeypatch): + from postgkyl.gpython.array import GkylArray + + def _fake_read_field(path): + return { + "cells": np.array([10]), + "lower": np.array([0.0]), + "upper": np.array([1.0]) + }, GkylArray.alloc(1, 5) # 5 != 10 + + monkeypatch.setattr(gpython.rio, "read_field", _fake_read_field) + r = GkylCReader(F1, ctx={}) + with pytest.raises(IOError, match="ghost-cell layout"): + r.load() + + +# ------------------------------------------------------------------- mapping +def test_adjust_for_ghost_cells_shrinks_and_extends_bounds(): + lower = np.array([0.0]) + upper = np.array([10.0]) + cells = np.array([10]) + lo, up, c = mapping.adjust_for_ghost_cells(lower, upper, cells, (8, )) + assert c[0] == 8 + dz = 1.0 # (10-0)/10 + assert lo[0] == pytest.approx(-1.0 * dz) + assert up[0] == pytest.approx(10.0 + 1.0 * dz) + + +# -------------------------------------------------------------------- writer +def test_write_derives_out_name_from_source_file(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + a = pg.load(F1).interpolate().select(comp=0) + a._file_name = "source.gkyl" + out = a.save() # out_name empty -> derived from _file_name + assert out == "source_mod.gkyl" or out.endswith("_mod.gkyl") + assert os.path.exists(out) + + +def test_write_appends_extension_when_missing(tmp_path): + a = pg.load(F1).interpolate().select(comp=0) + out = a.save(str(tmp_path / "no_ext"), extension="gkyl") + assert out.endswith("no_ext.gkyl") + assert os.path.exists(out) + + +def test_write_npy_and_txt_and_rejects_unknown_extension(tmp_path): + a = pg.load(F1).interpolate().select(comp=0) + npy_path = writer.save(a, out_name=str(tmp_path / "out.npy"), extension="npy") + assert os.path.exists(npy_path) + loaded = np.load(npy_path) + np.testing.assert_allclose(loaded, np.asarray(a.values).squeeze()) + + txt_path = writer.save(a, out_name=str(tmp_path / "out.txt"), extension="txt") + assert os.path.exists(txt_path) + with open(txt_path) as fh: + lines = fh.readlines() + assert len(lines) == int(np.prod(a.num_cells)) + + with pytest.raises(ValueError, match="Unsupported"): + writer.save(a, out_name=str(tmp_path / "out.bad"), extension="bad") + + +def test_write_txt_multidim_computes_row_major_strides(tmp_path): + """``_write_txt``'s stride computation (``basis[d] = prod(cells[d+1:])``) + only has a loop body for num_dims >= 2 -- a 1-D dataset skips it.""" + b = pg.load(F2D).interpolate().select(comp=0) + txt_path = writer.save(b, + out_name=str(tmp_path / "out2d.txt"), + extension="txt") + with open(txt_path) as fh: + lines = fh.readlines() + assert len(lines) == int(np.prod(b.num_cells)) + + +# --------------------------------------------------------- writer / metadata +def test_write_gkyl_roundtrips_metadata_through_meta_blob(tmp_path): + """DG poly order/basis type, physical params, and time/frame stamps read + off ``F1`` must survive a write() -> reload() round trip, not just the + raw field values.""" + a = pg.load(F1) + out = a.save(str(tmp_path / "roundtrip.gkyl"), extension="gkyl") + + reloaded = GkylReader(out, ctx={}) + reloaded.preload() + + for key in ("poly_order", "basis_type", "time", "frame", "changeset", + "builddate", "geometry_type", "Description"): + assert key in reloaded.ctx, f"{key!r} missing after round trip" + assert reloaded.ctx[key] == a.ctx[key] + + +def test_write_gkyl_roundtrips_custom_ctx_keys(tmp_path): + """Any ctx key that isn't structural/session-only (not just the keys + postgkyl special-cases) must be preserved verbatim.""" + a = pg.load(F1).interpolate().select(comp=0) + a.ctx["charge"] = -1.0 + a.ctx["mass"] = 1837.0 + out = a.save(str(tmp_path / "custom_meta.gkyl"), extension="gkyl") + + reloaded = GkylReader(out, ctx={}) + reloaded.preload() + assert reloaded.ctx["charge"] == -1.0 + assert reloaded.ctx["mass"] == 1837.0 + + +def test_write_gkyl_with_no_extra_ctx_writes_zero_meta_size(tmp_path): + """A dataset whose ctx carries only structural/session keys must produce + the same zero-length meta blob the writer always emitted -- no spurious + meta bytes for a dataset with nothing extra to say.""" + out_name = str(tmp_path / "no_meta.gkyl") + writer._write_gkyl(out_name, + num_dims=1, + num_comps=1, + num_cells=[4], + lo=[0.0], + up=[4.0], + values=np.arange(4, dtype=np.float64), + ctx={ + "cells": np.array([4]), + "lower": np.array([0.0]), + "upper": np.array([4.0]), + "grid_type": "uniform" + }) + + meta_size = np.fromfile(out_name, dtype=np.dtype("i8"), count=1, offset=21)[0] + assert meta_size == 0 + + reloaded = GkylReader(out_name, ctx={}) + reloaded.preload() + np.testing.assert_allclose(reloaded.cells, [4]) + + +def test_build_meta_excludes_internal_keys_and_renames_dg_fields(): + ctx = { + "cells": np.array([4]), + "lower": np.array([0.0]), + "upper": np.array([4.0]), + "num_comps": 1, + "num_dims": 1, + "grid_type": "uniform", + "value_form": "modal", + "num_quad": 3, + "interpolated": True, + "var_names": ["f"], + "poly_order": 2, + "basis_type": "serendipity", + "time": 0.5, + "frame": 3, + } + meta = writer._build_meta(ctx) + assert meta == { + "polyOrder": 2, + "basisType": "serendipity", + "time": 0.5, + "frame": 3 + } + + +def test_to_msgpack_safe_converts_numpy_scalars_and_arrays(): + assert writer._to_msgpack_safe(np.float64(1.5)) == 1.5 + assert isinstance(writer._to_msgpack_safe(np.float64(1.5)), float) + assert writer._to_msgpack_safe(np.int64(3)) == 3 + assert isinstance(writer._to_msgpack_safe(np.int64(3)), int) + assert writer._to_msgpack_safe(np.array([1.0, 2.0])) == [1.0, 2.0] + assert writer._to_msgpack_safe("serendipity") == "serendipity" + + +# --------------------------------------------------------------- gkyl_reader +def test_is_compatible_false_for_wrong_magic_and_missing_file(tmp_path): + bogus = tmp_path / "bad.gkyl" + bogus.write_bytes(b"definitely-not-gkyl-magic-bytes") + assert GkylReader(str(bogus), ctx={}).is_compatible() is False + assert GkylReader("/no/such/file.gkyl", ctx={}).is_compatible() is False + + +def test_defaults_ctx_to_a_fresh_dict_when_omitted(): + r = GkylReader(F1, ctx=None) + assert r.ctx == {"grid_type": "uniform"} + r.preload() + r.load() + + +def test_partial_load_negative_stop_component(): + r = GkylReader(F1, ctx={}, comp="0:-1") # drop the last component + r.preload() + grid, data = r.load() + assert data.shape[-1] == 5 + + +def test_partial_load_on_a_single_range_file_defaults_lo_up_idx(): + """``_get_data``'s ``lo_idx is None``/``up_idx is None`` defaults are only + reached for a file_type-1 (single-range) partial load -- type-3 multi-range + reads always pass explicit lo/up idx from the stored range headers.""" + full = GkylReader(F1D_SINGLE_RANGE, ctx={}) + full.preload() + _, full_data = full.load() + + r = GkylReader(F1D_SINGLE_RANGE, + ctx={}, + axes=("0:2", None, None, None, None, None)) + r.preload() + grid, data = r.load() + np.testing.assert_allclose(data, full_data[:2]) + + +@needs_gkeyll +def test_partial_load_excludes_a_whole_multirange_and_slices_axis(tmp_path): + """A real multi-range fixture (4 ranges of 6 cells); selecting exactly + range 0 forces the other 3 ranges' data blocks to be empty, exercising + the partial-load domain math, ``_get_block`` and the 'skip empty range' + continuation in ``_read_t3_v1_data``.""" + full = GkylReader(F1, ctx={}) + full.preload() + _, full_data = full.load() + + r = GkylReader(F1, ctx={}, axes=("0:6", None, None, None, None, None)) + r.preload() + grid, data = r.load() + assert data.shape == (6, 6) + assert grid[0].shape == (7, ) + np.testing.assert_allclose(data, full_data[:6]) + + +def test_partial_load_digit_axis_and_digit_component(tmp_path): + r = GkylReader(F1, ctx={}, axes=("2", None, None, None, None, None), comp="1") + r.preload() + grid, data = r.load() + assert data.shape == (1, 1) # one cell, one component + + +def test_partial_load_negative_stop_and_colon_component(): + r = GkylReader(F1, + ctx={}, + axes=("0:-2", None, None, None, None, None), + comp="0:3") + r.preload() + grid, data = r.load() + assert data.shape[0] == 24 - 2 + assert data.shape[-1] == 3 + + +@needs_gkeyll +def test_dynvec_single_chunk_round_trip_via_pure_python_reader(tmp_path): + from postgkyl.gpython import rio + path = str(tmp_path / "series.gkyl") + time = np.array([0.0, 0.5, 1.0]) + values = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + rio.write_dynvec(path, time, values) + + r = GkylReader(path, ctx={}) + r.preload() + grid, data = r.load() + np.testing.assert_allclose(grid[0], time) + np.testing.assert_allclose(data, values) + + +@needs_gkeyll +def test_dynvec_multi_chunk_continuation(tmp_path): + """Two dynvec writes concatenated back-to-back simulate the append + pattern Gkeyll uses for a running time series -- the reader must loop + back into ``_read_header`` for the second chunk without error.""" + from postgkyl.gpython import rio + p1, p2 = str(tmp_path / "c1.gkyl"), str(tmp_path / "c2.gkyl") + rio.write_dynvec(p1, np.array([0.0, 0.1]), np.array([[1.0, 2.0], [3.0, 4.0]])) + rio.write_dynvec(p2, np.array([0.2, 0.3, 0.4]), + np.array([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]])) + combo = tmp_path / "combo.gkyl" + combo.write_bytes(Path(p1).read_bytes() + Path(p2).read_bytes()) + + r = GkylReader(str(combo), ctx={}) + r.preload() + grid, data = r.load() + np.testing.assert_allclose(grid[0], [0.0, 0.1, 0.2, 0.3, 0.4]) + assert data.shape == (5, 2) + + +@needs_gkeyll +def test_dynvec_continuation_rejects_a_non_dynvec_second_chunk(tmp_path): + from postgkyl.gpython import rio + from postgkyl.gpython.array import GkylArray + p1 = str(tmp_path / "c1.gkyl") + rio.write_dynvec(p1, np.array([0.0, 0.1]), np.array([[1.0, 2.0], [3.0, 4.0]])) + pf = str(tmp_path / "field.gkyl") + rio.write_field(pf, { + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([3]) + }, GkylArray.from_numpy(np.ones((3, 2)))) + bad = tmp_path / "bad_combo.gkyl" + bad.write_bytes(Path(p1).read_bytes() + Path(pf).read_bytes()) + + r = GkylReader(str(bad), ctx={}) + r.preload() + with pytest.raises(TypeError, match="Inconsitent data"): + r.load() + + +def _write_legacy_v0_field(path, cells, lower, upper, data, real_type=2): + """Build a *version-0* raw field file: no gkyl0/version/type/meta header, + just real_type + the type-1 domain fields + data -- the format predating + the version-1 wrapper (see the module docstring in gkyl_reader.py).""" + dti = np.dtype("i8") + dtf = np.dtype("f4") if real_type == 1 else np.dtype("f8") + doffset = 4 if real_type == 1 else 8 + ndim = len(cells) + num_comps = data.shape[-1] + with open(path, "wb") as fh: + np.array([real_type], dtype=dti).tofile(fh) + np.array([ndim], dtype=dti).tofile(fh) + np.array(cells, dtype=dti).tofile(fh) + np.array(lower, dtype=dtf).tofile(fh) + np.array(upper, dtype=dtf).tofile(fh) + np.array([num_comps * doffset], dtype=dti).tofile(fh) + np.array([int(np.prod(cells))], dtype=dti).tofile(fh) + np.array(data, dtype=dtf).tofile(fh) + + +def test_legacy_version0_file_is_read_via_default_version_and_type(tmp_path): + path = str(tmp_path / "v0.gkyl") + data = np.arange(8, dtype=np.float64).reshape(4, 2) + _write_legacy_v0_field(path, [4], [0.0], [4.0], data) + + r = GkylReader(path, ctx={}) + assert r.is_compatible() is False # no "gkyl0" magic in this legacy format + r.preload() + grid, out = r.load() + assert r.version == 0 + np.testing.assert_allclose(grid[0], np.linspace(0.0, 4.0, 5)) + np.testing.assert_allclose(out, data) + + +def _write_v1_field(path, cells, lower, upper, data, real_type=2, meta=b""): + dti = np.dtype("i8") + dtf = np.dtype("f4") if real_type == 1 else np.dtype("f8") + doffset = 4 if real_type == 1 else 8 + ndim = len(cells) + num_comps = data.shape[-1] + with open(path, "wb") as fh: + np.array([103, 107, 121, 108, 48], dtype=np.dtype("b")).tofile(fh) + np.array([1], dtype=dti).tofile(fh) + np.array([1], dtype=dti).tofile(fh) + np.array([len(meta)], dtype=dti).tofile(fh) + fh.write(meta) + np.array([real_type], dtype=dti).tofile(fh) + np.array([ndim], dtype=dti).tofile(fh) + np.array(cells, dtype=dti).tofile(fh) + np.array(lower, dtype=dtf).tofile(fh) + np.array(upper, dtype=dtf).tofile(fh) + np.array([num_comps * doffset], dtype=dti).tofile(fh) + np.array([int(np.prod(cells))], dtype=dti).tofile(fh) + np.array(data, dtype=dtf).tofile(fh) + + +def test_single_precision_real_type_is_read_as_float32(tmp_path): + path = str(tmp_path / "f4.gkyl") + data = np.arange(6, dtype=np.float32).reshape(3, 2) + _write_v1_field(path, [3], [0.0], [3.0], data, real_type=1) + + r = GkylReader(path, ctx={}) + assert r.is_compatible() is True + r.preload() + grid, out = r.load() + assert r.dtf == np.dtype("f4") + np.testing.assert_allclose(out, data) + + +def test_reader_preserves_an_explicit_grid_type_and_accepts_axes_none(): + ctx = {"grid_type": "nodal"} + reader = GkylReader(F1D_SINGLE_RANGE, ctx=ctx, axes=None) + assert reader.ctx["grid_type"] == "nodal" + assert reader.partial_load is False + + +def test_non_mapping_metadata_is_ignored(tmp_path): + import msgpack + + path = str(tmp_path / "list-meta.gkyl") + data = np.arange(3, dtype=np.float64).reshape(3, 1) + _write_v1_field(path, [3], [0.0], [3.0], + data, + meta=msgpack.packb(["not", "a", "mapping"])) + reader = GkylReader(path, ctx={}) + reader.preload() + assert "basis_type" not in reader.ctx + + +def test_reader_metadata_overrides_are_independent(tmp_path): + path = str(tmp_path / "overrides.gkyl") + data = np.arange(3, dtype=np.float64).reshape(3, 1) + _write_v1_field(path, [3], [0.0], [3.0], data) + reader = GkylReader(path, + ctx={}, + basis_type="tensor", + poly_order=2, + value_form="nodal") + reader.preload() + assert reader.ctx["basis_type"] == "tensor" + assert reader.ctx["poly_order"] == 2 + assert reader.ctx["value_form"] == "nodal" + + +def test_full_slice_partial_load_uses_zero_offsets(): + reader = GkylReader(F1D_SINGLE_RANGE, + ctx={}, + axes=(":", None, None, None, None, None), + comp=":") + reader.preload() + _, values = reader.load() + assert values.shape == (8, 2) + + +def test_preload_and_field_load_allow_an_empty_context(tmp_path): + path = str(tmp_path / "empty-context.gkyl") + data = np.arange(3, dtype=np.float64).reshape(3, 1) + _write_v1_field(path, [3], [0.0], [3.0], data) + reader = GkylReader(path, ctx={}) + reader.ctx = {} + reader.preload() + assert reader.ctx == {} + _, values = reader.load() + np.testing.assert_allclose(values, data) + assert reader.ctx == {} + + +def test_dynvector_load_allows_an_empty_context(): + reader = GkylReader(DYNVEC, ctx={}) + reader.preload() + reader.ctx = {} + grid, values = reader.load() + assert len(grid[0]) == values.shape[0] + assert reader.ctx == {} + + +def test_load_raises_for_an_unsupported_file_type(tmp_path): + path = str(tmp_path / "v1.gkyl") + _write_v1_field(path, [3], [0.0], [3.0], np.zeros((3, 1))) + r = GkylReader(path, ctx={}) + r.preload() + r.file_type = 99 # not 1, 2, or 3; version is 1 so the version==0 branch + # doesn't rescue it either + with pytest.raises(TypeError, match="not presently supported"): + r.load() diff --git a/tests/test_coverage_leaf.py b/tests/test_coverage_leaf.py new file mode 100644 index 00000000..da71a663 --- /dev/null +++ b/tests/test_coverage_leaf.py @@ -0,0 +1,438 @@ +"""Coverage-completing tests for the leaf/engine/backend layers: numerics, +dg (interpolate/modal/rep), the remaining gpython corners (array/kernels), and the +matplotlib render backend. + +Run: PYTHONPATH=src pytest tests/test_coverage_leaf.py -v +""" + +import importlib +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg # noqa: E402 +from postgkyl import gpython, dg # noqa: E402 +# NB: `postgkyl.numerics.idx_parser` (the submodule) is shadowed by the +# `idx_parser` FUNCTION that numerics/__init__.py re-exports under the same +# attribute name -- both plain `from ... import idx_parser` and +# `import a.b.idx_parser as x` (itself sugar for `x = a.b.idx_parser`, an +# *attribute* lookup) resolve to the function. `importlib` sidesteps the +# package's __init__ entirely and returns the actual submodule object. +ip = importlib.import_module("postgkyl.numerics.idx_parser") +from postgkyl.numerics import elementwise # noqa: E402 +from postgkyl.gdatastate.gdatastate import GDataState # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +# ============================================================ numerics/idx_parser +def test_find_nearest_index_raises_without_a_coordinate_array(): + with pytest.raises(TypeError, match="no coordinate array"): + ip._find_nearest_index(None, 1.0) + + +def test_find_nearest_index_edge_cases(): + arr = np.array([0.0, 1.0, 2.0, 3.0]) + assert ip._find_nearest_index(arr, 10.0) == 2 # beyond the end -> idx-2 + assert ip._find_nearest_index(arr, -10.0) == 0 # before the start -> idx==0 + + +def test_find_cell_index_raises_without_a_coordinate_array(): + with pytest.raises(TypeError, match="no coordinate array"): + ip._find_cell_index(None, 1.0) + + +def test_string_to_index_rejects_non_strings(): + with pytest.raises(TypeError, match="not a string"): + ip._string_to_index(1.5, np.array([0.0, 1.0])) + + +def test_string_to_index_parses_a_float_string(): + arr = np.array([0.0, 1.0, 2.0, 3.0]) + assert ip._string_to_index("1.4", arr) == 1 + assert ip._string_to_index("1.4", arr, nodal=True) == 2 + + +def test_idx_parser_slice_with_empty_start_and_stop(): + arr = np.array([0.0, 1.0, 2.0, 3.0]) + s = ip.idx_parser("2:", arr) # empty stop -> len(array) + assert s == slice(2, 4) + s2 = ip.idx_parser(":2", arr) # empty start -> 0 + assert s2 == slice(0, 2) + + +def test_idx_parser_slice_negative_stop(): + arr = np.array([0.0, 1.0, 2.0, 3.0]) + assert ip.idx_parser("0:-1", arr) == slice(0, 4) + + +def test_idx_parser_slice_with_non_integer_stop_falls_back_to_float_lookup(): + """``hi`` failing int() parsing (a float-valued stop) is swallowed by the + ``except ValueError: pass`` guard, then resolved via the float-coordinate + path instead of the integer-count adjustment.""" + arr = np.array([0.0, 1.0, 2.0, 3.0]) + s = ip.idx_parser("0:1.4", arr) + assert s == slice(0, 1) + + +def test_idx_parser_rejects_unsupported_types(): + with pytest.raises(TypeError, match="Unsupported selector type"): + ip.idx_parser(3.0 + 4.0j) + + +# ============================================================ numerics/elementwise +def test_grids_compatible_rejects_different_ndims(): + a = [np.linspace(0.0, 1.0, 4)] + b = [np.linspace(0.0, 1.0, 4), np.linspace(0.0, 1.0, 4)] + assert elementwise.grids_compatible(a, b) is False + + +def test_grid_is_prefix_rejects_out_of_range_lengths(): + same_len = [np.linspace(0.0, 1.0, 4)] + assert elementwise.grid_is_prefix(same_len, + same_len) is False # not strictly smaller + assert elementwise.grid_is_prefix([], same_len) is False # empty + + +# ===================================================================== dg/interpolate +@needs_gkeyll +def test_interpolate_degenerates_1d_hybrid_to_serendipity(): + nb = dg.num_basis(1, 1, "serendipity") + values = np.zeros((5, nb)) + grid = [np.linspace(0.0, 1.0, 6)] + grid_out, out = dg.interpolate(values, + grid, + poly_order=1, + basis_type="hybrid") + assert out.shape[-1] == 1 + + +@needs_gkeyll +def test_interpolate_converts_nodal_basis_data_through_nodal_to_modal(): + """A NumPy-backed dataset tagged ``value_form="nodal"`` at load time + forces the nodal-basis-file convention -- exercising the nodal-to-modal + conversion machinery inside ``dg.interpolate`` (the values themselves are + meaningless here, only the code path and output shape matter). Native + (gkyl-backed) data always enforces the modal value_form, so this must be + plain NumPy from the start.""" + nb = dg.num_basis(1, 1, "serendipity") + d = pg.GData(ctx={ + "poly_order": 1, + "basis_type": "serendipity", + "value_form": "nodal" + }) + d.push([np.linspace(0.0, 1.0, 5)], np.zeros((4, nb))) + out = d.interpolate() + assert out.is_interpolated + assert out.values.ndim == 2 + + +@needs_gkeyll +def test_local_poly_degenerates_1d_hybrid_to_serendipity(): + nb = dg.num_basis(1, 1, "serendipity") + values = np.zeros((5, nb)) + grid = [np.linspace(0.0, 1.0, 6)] + grid_out, out = dg.local_poly(values, grid, poly_order=1, basis_type="hybrid") + assert out.shape[-1] == 1 + + +@needs_gkeyll +def test_local_poly_converts_nodal_basis_data_through_nodal_to_modal(): + nb = dg.num_basis(1, 1, "serendipity") + d = pg.GData(ctx={ + "poly_order": 1, + "basis_type": "serendipity", + "value_form": "nodal" + }) + d.push([np.linspace(0.0, 1.0, 5)], np.zeros((4, nb))) + out = d.local_poly() + assert out.is_interpolated + assert out.values.ndim == 2 + + +# ======================================================================= dg/modal +@needs_gkeyll +def test_modal_power_non_integer_exponent_uses_powsqrt(): + """A fractional exponent used to raise; it now routes through + ``dg.modal.powsqrt`` (``gkyl_proj_powsqrt_on_basis``) instead, field by + field, and must match the defining identity ``f ** 1.5 == f * (f ** 0.5)`` + (the latter computed via the same powsqrt kernel at a different exponent, + so this pins the exponent-doubling translation rather than just "some + value came out").""" + a = pg.load(F1) + cubed_half = a**1.5 + half = a**0.5 + expect = dg.modal.weak_mul(a.ctx["basis_type"], a.num_dims, + a.ctx["poly_order"], a.native, half.native) + np.testing.assert_allclose(cubed_half.native.view(), expect.view(), atol=1e-8) + + +@needs_gkeyll +def test_modal_power_without_cells_raises_for_non_integer_exponent(): + """``cells`` builds the powsqrt kernel's index range; calling ``dg.modal. + power`` directly (bypassing ``operations.arithmetic``, which always + supplies it from ``ctx["cells"]``) with a non-integer exponent and no + ``cells`` must fail clearly rather than crash inside the kernel.""" + a = _const_gkyl_array("serendipity", 1, 1, [4], 3.0) + with pytest.raises(ValueError, match="needs cells="): + dg.modal.power("serendipity", 1, 1, a, 0.5) + + +@needs_gkeyll +def test_modal_powsqrt_rejects_multi_field_ncomp_mismatch(): + """``a.ncomp`` must be a multiple of the basis's ``num_basis``.""" + a = gpython.GkylArray.alloc(3, 4) # 3 is not a multiple of num_basis (2) + with pytest.raises(ValueError, match="not a multiple"): + dg.modal.powsqrt("serendipity", 1, 1, [4], a, 1.0) + + +def _const_gkyl_array(basis_type, ndim, p, cells, value): + nb = gpython.basis.num_basis(basis_type, ndim, p) + b0 = 2.0**(-ndim / 2.0) + coeffs = np.zeros((int(np.prod(cells)), nb)) + coeffs[:, 0] = value / b0 + return gpython.GkylArray.from_numpy(coeffs) + + +@needs_gkeyll +def test_modal_average_full_reduction_corrects_the_raw_kernel_value(): + """Unlike the raw kernel (test_array_average_full_reduction_unweighted_ + writes_a_raw_value in test_gpython_kernels.py), ``dg.modal.average`` + rescales the degenerate (every dim averaged), unweighted case back into a + properly b0-normalized coefficient -- so it agrees with the weighted + case (which the underlying weak division already normalizes) and with + every other modal dataset's "value = coeff0 * b0" convention.""" + basis_type, p, cells = "serendipity", 1, [4] + a = _const_gkyl_array(basis_type, 1, p, cells, 3.0) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([2.0]), + "cells": np.array(cells) + } + keep_dirs, cells_avg, out = dg.modal.average(grid, basis_type, 1, p, a, [0]) + assert keep_dirs == [] + assert cells_avg == [1] + b0 = 2.0**(-1 / 2.0) + np.testing.assert_allclose(out.view()[0, 0] * b0, 3.0, atol=1e-10) + + +@needs_gkeyll +def test_modal_average_multi_field_loops_and_reassembles_per_field(): + """``gkyl_array_average`` has no field-index argument, so a multi-field + array (ncomp = nfields * num_basis) must be split, averaged one field at + a time, and reassembled -- verify each field's result matches averaging + it alone.""" + basis_type, p, cells = "serendipity", 1, [4] + nb = gpython.basis.num_basis(basis_type, 1, p) + values = [3.0, -1.5] + coeffs = np.concatenate( + [_const_gkyl_array(basis_type, 1, p, cells, v).view() for v in values], + axis=-1) + a = gpython.GkylArray.from_numpy(coeffs) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([2.0]), + "cells": np.array(cells) + } + keep_dirs, cells_avg, out = dg.modal.average(grid, basis_type, 1, p, a, [0]) + assert out.ncomp == 2 * nb + b0 = 2.0**(-1 / 2.0) + np.testing.assert_allclose(out.view()[0, 0] * b0, values[0], atol=1e-10) + np.testing.assert_allclose(out.view()[0, nb] * b0, values[1], atol=1e-10) + + +@needs_gkeyll +def test_modal_average_rejects_dirs_out_of_range(): + basis_type, p, cells = "serendipity", 1, [4] + a = _const_gkyl_array(basis_type, 1, p, cells, 1.0) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array(cells) + } + with pytest.raises(ValueError, match="out of range"): + dg.modal.average(grid, basis_type, 1, p, a, [1]) + + +@needs_gkeyll +def test_modal_average_rejects_ncomp_not_a_multiple_of_num_basis(): + a = gpython.GkylArray.alloc(3, 4) # num_basis for ser p1 1D is 2 + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([4]) + } + with pytest.raises(ValueError, match="not a multiple"): + dg.modal.average(grid, "serendipity", 1, 1, a, [0]) + + +@needs_gkeyll +def test_modal_average_rejects_weight_ncomp_mismatch(): + basis_type, p, cells = "serendipity", 1, [4] + a = _const_gkyl_array(basis_type, 1, p, cells, 1.0) + w = gpython.GkylArray.alloc(3, 4) # wrong ncomp for this basis + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array(cells) + } + with pytest.raises(ValueError, match="weight ncomp"): + dg.modal.average(grid, basis_type, 1, p, a, [0], weight=w) + + +# ==================================================================== gpython/array +@needs_gkeyll +def test_gkylarray_from_numpy_rejects_scalar_input(monkeypatch): + """``np.ascontiguousarray`` itself always promotes a 0-d input to 1-D, so + this guard can't be reached through any real ndarray -- it defends against + a hypothetical future NumPy behavior change. Drive it directly by faking + ascontiguousarray's return value.""" + from postgkyl.gpython import array as array_mod + monkeypatch.setattr(array_mod.np, + "ascontiguousarray", + lambda values, dtype=None: np.array(5.0, dtype=dtype)) + with pytest.raises(ValueError, match="at least a 1-D"): + gpython.GkylArray.from_numpy(np.array(5.0)) + + +# ==================================================================== gpython/kernels +@needs_gkeyll +def test_weak_mul_conf_phase_rejects_unsupported_phase_basis(): + from postgkyl.gpython import kernels as k + cop = gpython.GkylArray.alloc(2, 3) + pop = gpython.GkylArray.alloc(2, 12) + with pytest.raises(NotImplementedError, match="cross-mul supports"): + k.weak_mul_conf_phase("serendipity", 1, "bogus-basis", 2, 1, [3], [3, 4], + cop, pop) + + +@needs_gkeyll +def test_weak_mul_conf_phase_rejects_pop_ncomp_mismatch(): + from postgkyl.gpython import kernels as k + cbasis = gpython.basis.get_basis("serendipity", 1, 1) + pbasis = gpython.basis.get_basis("serendipity", 2, 1) + cop = gpython.GkylArray.alloc(cbasis.num_basis, 3) + pop = gpython.GkylArray.alloc(pbasis.num_basis + 1, 12) # wrong ncomp + with pytest.raises(ValueError, match="pop.ncomp"): + k.weak_mul_conf_phase("serendipity", 1, "serendipity", 2, 1, [3], [3, 4], + cop, pop) + + +# ======================================================================= dg/rep +@needs_gkeyll +def test_apply_per_field_rejects_ncomp_not_a_multiple(): + arr = gpython.GkylArray.alloc(3, 4) # ncomp=3, not a multiple of num_basis=2 + with pytest.raises(ValueError, match="not a multiple"): + dg.rep.modal_to_nodal("serendipity", 1, 1, arr) + + +@needs_gkeyll +def test_materialize_rejects_ncomp_not_a_multiple_of_points_per_cell(): + a = pg.load(F1) + arr = gpython.GkylArray.alloc(a.native.ncomp + 1, a.native.size) # off by one + with pytest.raises(ValueError, match="points/cell"): + dg.rep.materialize("serendipity", 1, 1, arr, a.grid, "nodal") + + +@needs_gkeyll +def test_tensor_point_layout_rejects_a_non_tensor_lin_index_collision( + monkeypatch): + """A hand-crafted node set whose per-dimension unique counts multiply to + ``num_basis`` (passing the coarse check) yet still contains a duplicate + cell -> point mapping (failing the fine-grained tensor-linearization + check): both are real defensive checks in ``_tensor_point_layout``, but + Gkeyll's actual basis node sets never exhibit either failure mode, so we + drive them directly by faking ``node_coords``.""" + from postgkyl.dg import rep + + duplicate_coords = np.array([[0., 0.], [0., 1.], [1., 0.], [0., 0.]]) + monkeypatch.setattr(rep.gpython_basis, "node_coords", + lambda *a, **k: duplicate_coords) + with pytest.raises(ValueError, match="not a tensor product"): + rep._tensor_point_layout("serendipity", 2, 1, "nodal", None) + + +@needs_gkeyll +def test_tensor_point_layout_rejects_misaligned_node_coordinates(monkeypatch): + from postgkyl.dg import rep + + nan_coords = np.array([[0.0], [np.nan]]) + monkeypatch.setattr(rep.gpython_basis, "node_coords", + lambda *a, **k: nan_coords) + with pytest.raises(ValueError, match="do not align on a tensor grid"): + rep._tensor_point_layout("serendipity", 1, 1, "nodal", None) + + +# =================================================================== render +@needs_gkeyll +def test_plot_rejects_empty_and_valueless_datasets(): + from postgkyl import render + with pytest.raises(ValueError, match="nothing to plot"): + render.plot() + + empty = GDataState() + with pytest.raises(ValueError, match="no values to plot"): + render.plot(empty) + + +@needs_gkeyll +def test_plot_multi_dataset_1d_with_labels_shows_legend_and_title(): + from postgkyl import render + a = pg.load(F1).interpolate().select(comp=0) + b = pg.load(F1).interpolate().select(comp=0) + fig = render.plot(a, + b, + multiblock=True, + legend_labels=["first", "second"], + title="my title", + no_show=True) + assert fig is not None + assert fig._suptitle is not None + assert fig._suptitle.get_text() == "my title" + + +@needs_gkeyll +def test_plot_rejects_more_than_two_dimensions(): + from postgkyl import render + d = GDataState() + d.push([np.linspace(0, 1, 3), + np.linspace(0, 1, 3), + np.linspace(0, 1, 3)], np.zeros((2, 2, 2, 1))) + with pytest.raises(ValueError, + match="Only 1D and 2D plots are currently supported"): + render.plot(d) + + +@needs_gkeyll +def test_plot_show_true_does_not_error_with_agg_backend(monkeypatch): + # matplotlib's own FigureCanvasBase.show() silently no-ops (no warning) on + # a genuinely headless Linux host (no DISPLAY) -- a deliberate + # headless-friendliness special case, not something this test is about. + # Force a DISPLAY so the assertion below exercises the actual behavior + # under test (Agg backend + no_show=False -> warn, don't raise) regardless of + # whether this host has a real display. + monkeypatch.setenv("DISPLAY", ":0") + a = pg.load(F1).interpolate().select(comp=0) + with pytest.warns(UserWarning, match="non-interactive"): + fig = a.plot(no_show=False) + assert fig is not None diff --git a/tests/test_coverage_operations.py b/tests/test_coverage_operations.py new file mode 100644 index 00000000..4ad12d39 --- /dev/null +++ b/tests/test_coverage_operations.py @@ -0,0 +1,612 @@ +"""Coverage-completing tests for the ``operations`` verb layer. + +The golden-path tests exercise ``comp=`` selection, the happy arithmetic +paths, and the default basis/poly_order. This file targets the error edges +and the less obvious dispatch branches: coordinate (``z0``) selection, +mixed-value_form/mixed-basis rejections, modal-scalar operator +combinations, ufunc edge cases, and every verb's metadata-missing guard. + +Run: PYTHONPATH=src pytest tests/test_coverage_ops.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg # noqa: E402 +from postgkyl import gpython, operations # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F3 = os.path.join(DATA, "generated", "3d_ms_p1.gkyl") + + +def _dynvec_dataset(tmp_path, time, values): + from postgkyl.gpython import rio + path = str(tmp_path / "series.gkyl") + rio.write_dynvec(path, np.asarray(time), np.asarray(values)) + return pg.load(path) + + +# ============================================================== operations.select +@needs_gkeyll +def test_select_by_coordinate_on_a_nodal_grid(tmp_path): + """A dynvector's grid length equals its value count exactly, so + ``select``'s ``is_matching`` branch is True (unlike interpolated field + data, whose grid is always one edge longer than its values).""" + d = _dynvec_dataset(tmp_path, [0.0, 0.5, 1.0, 1.5], + [[1.0], [2.0], [3.0], [4.0]]) + + by_int = d.select(z0=1) + np.testing.assert_allclose(by_int.values, [[2.0]]) + np.testing.assert_allclose(by_int.grid[0], [0.5]) + + by_float = d.select(z0=0.6) + np.testing.assert_allclose(by_float.values, [[3.0]]) + + by_slice = d.select(z0="1:3") + np.testing.assert_allclose(by_slice.values, [[2.0], [3.0]]) + np.testing.assert_allclose(by_slice.grid[0], [0.5, 1.0]) + + by_negative_int = d.select(z0=-1) + np.testing.assert_allclose(by_negative_int.values, [[4.0]]) + + with pytest.raises(TypeError, match="single index or a slice"): + d.select( + z0="1,2") # comma selector is comp-only syntax, not valid for z-axes + + +@needs_gkeyll +def test_select_by_coordinate_on_a_non_matching_edge_grid(): + """Interpolated field data: the grid has one more point than the values + along every axis (edges vs. cell values) -- the ``is_matching`` False + path.""" + g = pg.load(F1).interpolate() + assert g.grid[0].shape[0] == g.values.shape[0] + 1 + + by_float = g.select(z0=0.0) + assert by_float.values.shape[0] == 1 + by_slice = g.select(z0="2:5") + assert by_slice.values.shape[0] == 3 + assert by_slice.grid[0].shape[0] == 4 + + +@needs_gkeyll +def test_select_keeps_native_point_values_in_the_native_backend(): + nodal = pg.load(F1).to_nodal() + selected = nodal.select(comp=0) + assert selected.backend == "gkyl" + assert selected.ctx["value_form"] == "nodal" + np.testing.assert_array_equal(selected.ctx["cells"], [24]) + assert selected.native.ncomp == 1 + + +# ========================================================== operations.arithmetic +@needs_gkeyll +def test_numpy_domain_rejects_incompatible_grids_and_shapes(): + a = pg.load(F1).interpolate() + b = pg.load(F1).interpolate() + b_sub = b.select(comp=0) # different shape than the full 'a' + with pytest.raises(ValueError, match="incompatible shapes"): + a + b_sub + + c = pg.load(F1).interpolate() + c.grid[0] = c.grid[0] + 1.0 # displace the grid -> no longer "compatible" + with pytest.raises(ValueError, match="different grids"): + a + c + + +@needs_gkeyll +def test_basis_of_raises_when_metadata_missing(): + a, b = pg.load(F1), pg.load(F1) + del a.ctx["poly_order"] + with pytest.raises(ValueError, match="basis_type/poly_order"): + a * b + + +@needs_gkeyll +def test_modal_binary_rejects_mixing_with_a_plain_array(): + a = pg.load(F1) + with pytest.raises(ValueError, match="cannot mix native modal data"): + a * np.zeros((24, 6)) + + +@needs_gkeyll +def test_modal_dataset_pair_rejects_grid_and_basis_mismatch(): + a = pg.load(F1) + b = pg.load(F1) + b.grid[0] = b.grid[0] + 100.0 + with pytest.raises(ValueError, match="different grids"): + a * b + + c = pg.load(F1) + c.ctx["basis_type"] = "tensor" + with pytest.raises(ValueError, match="different DG bases"): + a * c + + +@needs_gkeyll +def test_modal_dataset_pair_rejects_unsupported_op(): + a, b = pg.load(F1), pg.load(F1) + with pytest.raises(ValueError, match="not defined between two"): + a**b + + +@needs_gkeyll +def test_conf_phase_mul_requires_both_operands_modal(): + """Mixed value_form on a conf*phase multiply (different num_dims) + must refuse just like the same-dims path, not silently coerce.""" + conf_edges = [np.linspace(0.0, 1.0, 4)] + phase_edges = [np.linspace(0.0, 1.0, 4), np.linspace(-1.0, 1.0, 5)] + cbasis = gpython.basis.get_basis("serendipity", 1, 1) + pbasis = gpython.basis.get_basis("hybrid", 2, 1) + + conf = pg.GData() + conf.ctx.update(basis_type="serendipity", + poly_order=1, + value_form="modal", + cells=np.array([3])) + conf.push(conf_edges, + gpython.array.GkylArray.from_numpy(np.zeros((3, cbasis.num_basis)))) + + phase = pg.GData() + phase.ctx.update(basis_type="hybrid", + poly_order=1, + value_form="modal", + cells=np.array([3, 4])) + phase.push( + phase_edges, + gpython.array.GkylArray.from_numpy(np.zeros((12, pbasis.num_basis)))) + + phase_nodal = phase.to_nodal() + with pytest.raises(ValueError, match="modal DG coefficients only"): + conf * phase_nodal + + +@needs_gkeyll +@pytest.mark.parametrize( + "expr", + [ + lambda a: a / 2.0, # modal / scalar (linear divide) + lambda a: 5.0 - a, # scalar - modal + lambda a: a - 5.0, # modal - scalar + lambda a: 5.0 / a, # scalar / modal (weak reciprocal) + ]) +def test_modal_scalar_operator_combinations(expr): + a = pg.load(F1) + out = expr(a) + assert isinstance(out, pg.GData) + assert out.backend == "gkyl" + + +@needs_gkeyll +def test_modal_scalar_rejects_reflected_power(): + a = pg.load(F1) + with pytest.raises(ValueError, match="not defined for modal"): + 2.0**a + + +@needs_gkeyll +def test_apply_ufunc_non_reduction_method_and_out_kwarg_are_rejected(): + a = pg.load(F1).interpolate() + assert a.__array_ufunc__(np.add, "accumulate", a) is NotImplemented + assert a.__array_ufunc__(np.sqrt, "__call__", a, + out=(np.zeros(1), )) is NotImplemented + + +def test_apply_ufunc_reductions_return_numpy_results(): + values = np.array([[3.0, 2.0], [-4.0, 5.0], [1.0, -2.0]]) + a = pg.GData() + a.push([np.linspace(0.0, 1.0, 4)], values) + + assert np.max(a) == np.max(values) + assert np.min(a) == np.min(values) + assert np.sum(a) == np.sum(values) + assert np.prod(a) == np.prod(values) + np.testing.assert_allclose(np.max(a, axis=0), np.max(values, axis=0)) + np.testing.assert_allclose( + np.sum(a, axis=1, keepdims=True, dtype=np.float64), + np.sum(values, axis=1, keepdims=True, dtype=np.float64)) + + bool_values = values > 0 + b = pg.GData() + b.push([np.linspace(0.0, 1.0, 4)], bool_values) + assert np.all(b) == np.all(bool_values) + assert np.any(b) == np.any(bool_values) + + +def test_remaining_reflected_and_unary_operators_preserve_data(): + a = pg.GData() + a.push([np.linspace(0.0, 1.0, 4)], np.array([[-2.0], [0.0], [3.0]])) + + np.testing.assert_allclose((1.0 + a).values, 1.0 + a.values) + np.testing.assert_allclose(abs(a).values, np.abs(a.values)) + positive = +a + np.testing.assert_allclose(positive.values, a.values) + assert positive is not a + + +def test_apply_ufunc_reduction_rejects_non_dataset_inputs(): + assert operations.arithmetic.apply_ufunc(np.add, "reduce", 1.0, + 2.0) is NotImplemented + + +@needs_gkeyll +def test_apply_ufunc_rejects_shape_mismatch(): + a = pg.load(F1).interpolate() + b = pg.load(F1).interpolate().select(comp=0) + with pytest.raises(ValueError, match="incompatible shapes"): + np.add(a, b) + + +@needs_gkeyll +def test_apply_ufunc_accepts_scalars_and_rejects_unhandled_types(): + a = pg.load(F1).interpolate() + out = np.add(a, 2.0) + np.testing.assert_allclose(out.values, a.values + 2.0) + assert a.__array_ufunc__(np.add, "__call__", a, + "not-a-number") is NotImplemented + + +# ========================================================== operations.interpolate +def test_interpolate_requires_basis_type_when_none_given(): + d = pg.GData() + d.push([np.linspace(0.0, 1.0, 4)], np.zeros((3, 2))) + with pytest.raises(ValueError, match="no 'basis_type' metadata"): + d.interpolate() + + +def test_interpolate_requires_poly_order_when_none_given(): + d = pg.GData() + d.ctx["basis_type"] = "serendipity" + d.push([np.linspace(0.0, 1.0, 4)], np.zeros((3, 2))) + with pytest.raises(ValueError, match="no 'poly_order' metadata"): + d.interpolate() + + +# =========================================================== operations.represent +@needs_gkeyll +def test_represent_rejects_numpy_backed_and_missing_metadata(): + interpolated = pg.load(F1).interpolate() + with pytest.raises(ValueError, match="NumPy-backed"): + interpolated.to_modal() + + a = pg.load(F1) + del a.ctx["poly_order"] + with pytest.raises(ValueError, match="no basis_type/poly_order"): + a.to_nodal() + + +@needs_gkeyll +def test_represent_rejects_unknown_target(): + a = pg.load(F1) + with pytest.raises(ValueError, match="unknown value_form"): + operations.represent(a, to="bogus") + + +@needs_gkeyll +def test_represent_rejects_quad_dataset_missing_num_quad(): + q = pg.load(F1).to_quad() + del q.ctx["num_quad"] + with pytest.raises(ValueError, match="lost its 'num_quad'"): + q.to_modal() + + +@needs_gkeyll +def test_represent_same_representation_clones(): + a = pg.load(F1) + same = a.to_modal() # already modal -> the "cur == to" clone branch + np.testing.assert_allclose(same.values, a.values) + assert same.native is not a.native + + +@needs_gkeyll +def test_apply_rejects_non_modal_data(): + a = pg.load(F1).to_nodal() + with pytest.raises(ValueError, match="expects modal data"): + a.apply(np.sqrt) + + +# =============================================================== operations.info +@needs_gkeyll +def test_info_verb_handles_multiple_datasets(): + a, b = pg.load(F1), pg.load(F1) + summaries = pg.info(a, b) + assert len(summaries) == 2 + assert all("Number of components" in s for s in summaries) + + +# ============================================================ operations.integrate +@needs_gkeyll +def test_integrate_requires_basis_metadata(): + a = pg.load(F1) + del a.ctx["basis_type"] + with pytest.raises(ValueError, match="basis_type/poly_order"): + a.integrate() + + +# ============================================================== operations.average +@needs_gkeyll +def test_average_full_reduction_matches_integrate_over_volume(): + a = pg.load(F1) + avg = a.average([0]) + assert avg.num_dims == 1 + assert avg.ctx["cells"].tolist() == [1] + lo, up = a.bounds + volume = float(up[0] - lo[0]) + integral = a.integrate() + b0 = 2.0**(-avg.num_dims / 2.0) + np.testing.assert_allclose(np.asarray(avg.native.view())[0, ::2] * b0, + np.asarray(integral) / volume, + rtol=1e-8) + + +@needs_gkeyll +def test_average_partial_reduction_of_a_constant_field(): + basis_type, p = "serendipity", 1 + cells = [4, 3] + nb = gpython.basis.num_basis(basis_type, 2, p) + b0 = 2.0**(-2 / 2.0) + coeffs = np.zeros((int(np.prod(cells)), nb)) + coeffs[:, 0] = 3.0 / b0 + + d = pg.GData() + d.ctx.update(basis_type=basis_type, + poly_order=p, + cells=np.array(cells), + value_form="modal") + grid = [ + np.linspace(0.0, 2.0, cells[0] + 1), + np.linspace(0.0, 1.0, cells[1] + 1) + ] + d.push(grid, gpython.array.GkylArray.from_numpy(coeffs)) + + out = d.average([1]) + assert out.num_dims == 1 + assert out.ctx["cells"].tolist() == [cells[0]] + assert out.grid[0].shape[0] == cells[0] + 1 + b0_avg = 2.0**(-1 / 2.0) + np.testing.assert_allclose(np.asarray(out.native.view())[:, 0] * b0_avg, + 3.0, + atol=1e-10) + + +@needs_gkeyll +def test_average_accepts_a_compatible_weight(): + basis_type, poly_order, cells = "serendipity", 1, [4] + num_basis = gpython.basis.num_basis(basis_type, 1, poly_order) + basis_constant = 2.0**0.5 + + def constant_state(value): + coefficients = np.zeros((cells[0], num_basis)) + coefficients[:, 0] = value * basis_constant + data = pg.GData( + ctx={ + "basis_type": basis_type, + "poly_order": poly_order, + "value_form": "modal", + "cells": np.array(cells), + }) + data.push([np.linspace(0.0, 1.0, cells[0] + 1)], + gpython.GkylArray.from_numpy(coefficients)) + return data + + out = constant_state(3.0).average([0], weight=constant_state(2.0)) + np.testing.assert_allclose(out.native.view()[0, 0] / basis_constant, + 3.0, + atol=1e-10) + + +@needs_gkeyll +def test_average_rejects_numpy_backed_and_non_modal(): + interpolated = pg.load(F1).interpolate() + with pytest.raises(ValueError, match="native modal data"): + interpolated.average([0]) + + nodal = pg.load(F1).to_nodal() + with pytest.raises(ValueError, match="modal value_form"): + nodal.average([0]) + + +@needs_gkeyll +def test_average_rejects_missing_basis_metadata(): + a = pg.load(F1) + del a.ctx["poly_order"] + with pytest.raises(ValueError, match="basis_type/poly_order"): + a.average([0]) + + +@needs_gkeyll +def test_average_rejects_weight_mismatch(): + a = pg.load(F1) + weight_wrong_dims = pg.GData() + weight_wrong_dims.ctx.update(basis_type="serendipity", + poly_order=1, + cells=np.array([4, 3]), + value_form="modal") + weight_wrong_dims.push([np.linspace(0.0, 1.0, 5), + np.linspace(0.0, 1.0, 4)], + gpython.array.GkylArray.from_numpy(np.zeros((12, 4)))) + with pytest.raises(ValueError, match="dims but the field has"): + a.average([0], weight=weight_wrong_dims) + + weight_wrong_basis = pg.load(F1) + weight_wrong_basis.ctx["basis_type"] = "tensor" + with pytest.raises(ValueError, match="basis_type"): + a.average([0], weight=weight_wrong_basis) + + weight_wrong_p = pg.load(F1) + weight_wrong_p.ctx["poly_order"] = 2 + with pytest.raises(ValueError, match="poly_order"): + a.average([0], weight=weight_wrong_p) + + +@needs_gkeyll +def test_average_tag_and_label_and_inplace(): + a = pg.load(F1) + out = a.average([0], tag="reduced", label="my label") + assert out.tag == "reduced" + assert out.label == "my label" + assert a.num_dims == 1 and a.ctx["cells"].tolist() != [1 + ] # original untouched + + b = pg.load(F1) + mutated = b.average([0], inplace=True) + assert mutated is b + assert b.ctx["cells"].tolist() == [1] + + +# ============================================================= operations.integrate +@needs_gkeyll +def test_integrate_partial_modal_stays_native_and_exact(): + a = pg.load(F3) + reduced = a.integrate(2) + assert reduced.backend == "gkyl" + assert reduced.ctx["value_form"] == "modal" + assert reduced.num_dims == 2 + np.testing.assert_allclose(reduced.integrate(), a.integrate(), rtol=1e-12) + + +def test_integrate_partial_point_data_removes_the_axis(): + a = pg.load(F3).interpolate() + r = a.integrate(2) + assert r.num_dims == 2 + assert r.num_cells.tolist() == list(a.num_cells[:2]) + assert r.ctx.get("interpolated") is True + + +def test_integrate_partial_point_data_matches_manual_sum(): + a = pg.load(F3).interpolate() + r = a.integrate(2) + dz = np.diff(a.grid[2]) + expected = np.tensordot(np.asarray(a.values), dz, axes=([2], [0])) + np.testing.assert_allclose(np.asarray(r.values), expected) + + +def test_integrate_point_default_is_a_full_terminal_integral(): + a = pg.load(F1).interpolate() + result = a.integrate() + assert isinstance(result, np.ndarray) + assert result.shape == (a.num_comps, ) + + +@needs_gkeyll +def test_integrate_partial_on_native_nodal_representation(): + # A gkyl-native nodal/quad dataset materializes to its true point grid + # before integrating -- same bridge ``plot`` uses (Doctrine V: one home). + nodal = pg.load(F3).to_nodal() + r = nodal.integrate(2) + assert r.backend == "numpy" + assert r.num_dims == 2 + assert r.ctx.get("value_form") is None # stale tag cleared, not "nodal" + + +def test_integrate_partial_tag_and_label(): + a = pg.load(F3).interpolate() + r = a.integrate(2, tag="reduced", label="my label") + assert r.tag == "reduced" + assert r.label == "my label" + assert a.num_dims == 3 # original left untouched (inplace=False default) + + +def test_integrate_partial_inplace_mutates_the_dataset(): + a = pg.load(F3).interpolate() + out = a.integrate(2, inplace=True) + assert out is a + assert a.num_dims == 2 + + +def test_integrate_full_rejects_partial_result_options(): + a = pg.load(F1).interpolate() + with pytest.raises(ValueError, match="partial integration"): + a.integrate(tag="not-a-dataset") + + +@pytest.mark.parametrize(("axis", "message"), [ + ((), "at least one axis"), + ((0, 0), "must be distinct"), + ((1, ), "out of range"), +]) +def test_integrate_rejects_invalid_axis_sets(axis, message): + a = pg.GData() + a.push([np.linspace(0.0, 1.0, 5)], np.ones((4, 1))) + with pytest.raises(ValueError, match=message): + a.integrate(axis) + + +def test_native_integration_guard_reports_backend_before_basis(): + from importlib import import_module + integrate_module = import_module("postgkyl.operations.integrate") + a = pg.GData() + a.push([np.linspace(0.0, 1.0, 5)], np.ones((4, 1))) + with pytest.raises(ValueError, match="needs native modal data"): + integrate_module._native_basis(a) + + +@needs_gkeyll +def test_native_integration_guard_rejects_point_value_forms(): + from importlib import import_module + integrate_module = import_module("postgkyl.operations.integrate") + nodal = pg.load(F1).to_nodal() + with pytest.raises(ValueError, match="expects the modal value_form"): + integrate_module._native_basis(nodal) + + +def test_integrate_rejects_native_only_op_for_point_data(): + one_dim = pg.GData() + one_dim.push([np.linspace(0.0, 1.0, 5)], np.ones((4, 1))) + with pytest.raises(ValueError, match="full native-DG"): + one_dim.integrate(op="abs") + + two_dim = pg.GData() + two_dim.push([np.linspace(0.0, 1.0, 5), + np.linspace(0.0, 1.0, 4)], np.ones((4, 3, 1))) + with pytest.raises(ValueError, match="full native-DG"): + two_dim.integrate(1, op="abs") + + +def test_remaining_mapped_axes_reindexes_surviving_groups(): + from importlib import import_module + integrate_module = import_module("postgkyl.operations.integrate") + data = pg.GData(ctx={"mapped_axes": {0: 0, 1: 0, 2: 2}}) + assert integrate_module._remaining_mapped_axes(data, [0, 1]) == {0: 0, 1: 0} + + +def test_curvilinear_lookup_ignores_flat_axes_and_can_miss_an_axis(): + from postgkyl.operations._curvilinear import block_for_axis, curvilinear_blocks + + blocks = curvilinear_blocks([np.arange(3), np.ones((3, 4))], { + 0: 0, + 1: 1, + }) + assert blocks == {1: [1]} + assert block_for_axis(blocks, 0) is None + + +def test_fft_preserves_an_already_point_aligned_grid(): + data = pg.GData() + data.push([np.linspace(0.0, 1.0, 8, endpoint=False)], + np.arange(8, dtype=float)[:, None]) + out = data.fft() + assert out.values.shape == data.values.shape + + +def test_val2coord_range_accepts_negative_slice_endpoints(): + from postgkyl.operations.val2coord import _get_range + + np.testing.assert_array_equal(_get_range("-4:-1", 6), [2, 3, 4]) + np.testing.assert_array_equal(_get_range("1:4", 6), [1, 2, 3]) diff --git a/tests/test_data/bimaxwellian-elc.gkyl b/tests/test_data/bimaxwellian-elc.gkyl deleted file mode 100644 index 2ad9e175..00000000 Binary files a/tests/test_data/bimaxwellian-elc.gkyl and /dev/null differ diff --git a/tests/test_data/bimaxwellian-jacobvel.gkyl b/tests/test_data/bimaxwellian-jacobvel.gkyl deleted file mode 100644 index aa6f044a..00000000 Binary files a/tests/test_data/bimaxwellian-jacobvel.gkyl and /dev/null differ diff --git a/tests/test_data/bimaxwellian-mapc2p-vel.gkyl b/tests/test_data/bimaxwellian-mapc2p-vel.gkyl deleted file mode 100644 index e9fe320e..00000000 Binary files a/tests/test_data/bimaxwellian-mapc2p-vel.gkyl and /dev/null differ diff --git a/tests/test_data/hll-euler.gkyl b/tests/test_data/hll-euler.gkyl deleted file mode 100644 index aac2eac1..00000000 Binary files a/tests/test_data/hll-euler.gkyl and /dev/null differ diff --git a/tests/test_data/shock-f-ser-p1.gkyl b/tests/test_data/shock-f-ser-p1.gkyl deleted file mode 100644 index 8d7d9903..00000000 Binary files a/tests/test_data/shock-f-ser-p1.gkyl and /dev/null differ diff --git a/tests/test_data/shock-f-ten-p1.gkyl b/tests/test_data/shock-f-ten-p1.gkyl deleted file mode 100644 index 8d7d9903..00000000 Binary files a/tests/test_data/shock-f-ten-p1.gkyl and /dev/null differ diff --git a/tests/test_data/shock-rtheta-ser.gkyl b/tests/test_data/shock-rtheta-ser.gkyl deleted file mode 100644 index 6f2db732..00000000 Binary files a/tests/test_data/shock-rtheta-ser.gkyl and /dev/null differ diff --git a/tests/test_data/shock-rtheta-ten.gkyl b/tests/test_data/shock-rtheta-ten.gkyl deleted file mode 100644 index 6f2db732..00000000 Binary files a/tests/test_data/shock-rtheta-ten.gkyl and /dev/null differ diff --git a/tests/test_data/twostream-f-p1.bp/data.0 b/tests/test_data/twostream-f-p1.bp/data.0 deleted file mode 100644 index 761bde7c..00000000 Binary files a/tests/test_data/twostream-f-p1.bp/data.0 and /dev/null differ diff --git a/tests/test_data/twostream-f-p1.bp/md.0 b/tests/test_data/twostream-f-p1.bp/md.0 deleted file mode 100644 index 66df4a7b..00000000 Binary files a/tests/test_data/twostream-f-p1.bp/md.0 and /dev/null differ diff --git a/tests/test_data/twostream-f-p1.bp/md.idx b/tests/test_data/twostream-f-p1.bp/md.idx deleted file mode 100644 index 0e2ccb96..00000000 Binary files a/tests/test_data/twostream-f-p1.bp/md.idx and /dev/null differ diff --git a/tests/test_data/twostream-f-p1.bp/mmd.0 b/tests/test_data/twostream-f-p1.bp/mmd.0 deleted file mode 100644 index f4b63f2a..00000000 Binary files a/tests/test_data/twostream-f-p1.bp/mmd.0 and /dev/null differ diff --git a/tests/test_data/twostream-f-p1.bp/profiling.json b/tests/test_data/twostream-f-p1.bp/profiling.json deleted file mode 100644 index 38c7e5ab..00000000 --- a/tests/test_data/twostream-f-p1.bp/profiling.json +++ /dev/null @@ -1,3 +0,0 @@ -[ -{ "rank":0, "start":"Wed_Aug_28_12:13:01_2024", "bytes":0, "AWD":{"mus":137, "nCalls":1}, "close_ts":{"mus":101, "nCalls":1}, "meta_lvl1":{"mus":3, "nCalls":1}, "meta_lvl2":{"mus":25, "nCalls":1}, "endstep":{"mus":269, "nCalls":1}, "transport_0":{"type":"File_POSIX", "close":{"mus":2, "nCalls":1}, "write":{"mus":41, "nCalls":1}, "open":{"mus":19, "nCalls":1}}, "transport_1":{"type":"File_POSIX", "close":{"mus":0, "nCalls":1}, "write":{"mus":1, "nCalls":5}, "open":{"mus":32, "nCalls":1}} } -] diff --git a/tests/test_data/twostream-f-p2.gkyl b/tests/test_data/twostream-f-p2.gkyl deleted file mode 100644 index 90fde841..00000000 Binary files a/tests/test_data/twostream-f-p2.gkyl and /dev/null differ diff --git a/tests/test_data/twostream-f-p2_0.bp b/tests/test_data/twostream-f-p2_0.bp deleted file mode 100644 index cb640c69..00000000 Binary files a/tests/test_data/twostream-f-p2_0.bp and /dev/null differ diff --git a/tests/test_data/twostream-f-p2_1.bp b/tests/test_data/twostream-f-p2_1.bp deleted file mode 100644 index 4650e838..00000000 Binary files a/tests/test_data/twostream-f-p2_1.bp and /dev/null differ diff --git a/tests/test_data/twostream-field-energy.bp b/tests/test_data/twostream-field-energy.bp deleted file mode 100644 index 8a1f22c6..00000000 Binary files a/tests/test_data/twostream-field-energy.bp and /dev/null differ diff --git a/tests/test_data/twostream-field-energy.gkyl b/tests/test_data/twostream-field-energy.gkyl deleted file mode 100644 index 8af2a893..00000000 Binary files a/tests/test_data/twostream-field-energy.gkyl and /dev/null differ diff --git a/tests/test_dg_differentiate_and_eval_at_coord_proj.py b/tests/test_dg_differentiate_and_eval_at_coord_proj.py new file mode 100644 index 00000000..e4538a1c --- /dev/null +++ b/tests/test_dg_differentiate_and_eval_at_coord_proj.py @@ -0,0 +1,246 @@ +"""Correctness tests for the ``differentiate``/``eval_at_coord_proj`` DG +operations ported from the old ``GkeyllDGops`` (ctypes) class -- now backed +by the compiled shim (``gpython.kernels``), orchestrated in ``dg.modal``, and +exposed as the ``eval_at_coord_proj`` verb/CLI command (``differentiate`` has +no dedicated verb: it is used internally, the way the old class's +``differentiate`` method fed gk-quantity math, not as a user-facing CLI verb). + +Run: PYTHONPATH=src pytest tests/test_dg_differentiate_and_eval_at_coord_proj.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.abspath(__file__)) +SRC = os.path.join(os.path.dirname(ROOT), "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import postgkyl as pg # noqa: E402 +from postgkyl import dg, gpython # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +DATA = os.path.join(ROOT, "test_data") +GKHYB = os.path.join(DATA, "rt_gk_tcv_iwl_1x2v_p1-elc_250.gkyl") + + +# ============================================================ differentiate +@needs_gkeyll +def test_differentiate_first_order_on_a_linear_field(): + """f(x) = x on a single cell covering physical x in [0, 2] (dx=2); df/dx + should be the constant field 1 everywhere.""" + # f(z) = c0*b0 + c1*b1 = c0/sqrt(2) + c1*sqrt(1.5)*z; matching f(z) = 1+z. + c0 = np.sqrt(2.0) + c1 = 1.0 / np.sqrt(1.5) + a = gpython.GkylArray.from_numpy(np.array([[c0, c1]])) + + out = dg.modal.differentiate("serendipity", + 1, + 1, + a, + dir=0, + diff_order=1, + dx=2.0) + + # A constant field 1 has coefficient c0' = sqrt(2), c1' = 0. + np.testing.assert_allclose(out.view(), [[np.sqrt(2.0), 0.0]], atol=1e-12) + + +@needs_gkeyll +def test_differentiate_second_order_on_a_quadratic_field(): + """f(z) = z^2 (p2 serendipity, single cell, dx=2 so d/dx = d/dz); + d^2f/dz^2 = 2 exactly, independent of position.""" + # Reference basis for 1x p2 serendipity: b0=1/sqrt(2), b1=sqrt(1.5)*z, + # b2 = sqrt(5/8)*(3z^2-1). f(z)=z^2 => project onto b2 (mean-subtracted + # quadratic) plus a constant: z^2 = 1/3 + (1/3)*(3z^2-1). + # c0/sqrt(2) = 1/3 => c0 = sqrt(2)/3; c2*sqrt(5/8) = 1/3 => c2 = 1/(3*sqrt(5/8)). + c0 = np.sqrt(2.0) / 3.0 + c2 = 1.0 / (3.0 * np.sqrt(5.0 / 8.0)) + a = gpython.GkylArray.from_numpy(np.array([[c0, 0.0, c2]])) + + out = dg.modal.differentiate("serendipity", + 1, + 2, + a, + dir=0, + diff_order=2, + dx=2.0) + + expected_c0 = 2.0 * np.sqrt(2.0) # constant field "2" -> coeff0 = 2*sqrt(2) + np.testing.assert_allclose(out.view()[0, 0], expected_c0, atol=1e-10) + np.testing.assert_allclose(out.view()[0, 1:], [0.0, 0.0], atol=1e-10) + + +@needs_gkeyll +def test_differentiate_rejects_out_of_table_combinations(): + a = gpython.GkylArray.alloc(4, 1) # 1x p1 serendipity-shaped, but wrong basis + with pytest.raises(NotImplementedError, match="serendipity/tensor"): + dg.modal.differentiate("gkhybrid", 2, 1, a, dir=0, diff_order=1, dx=1.0) + + a3 = gpython.GkylArray.alloc(gpython.basis.num_basis("tensor", 3, 1), 1) + with pytest.raises(NotImplementedError, match="ndim"): + dg.modal.differentiate("tensor", 3, 1, a3, dir=0, diff_order=1, dx=1.0) + + +# ======================================================== eval_at_coord_proj +@needs_gkeyll +def test_eval_at_coord_proj_matches_direct_polynomial_evaluation(): + """Cross-check: the target's reconstructed value at an arbitrary surviving + point must equal the donor's own reconstruction at the same point (with + the eliminated coordinate substituted) -- exactly, since both sides + evaluate the same underlying polynomial.""" + ndim, poly_order = 2, 1 + basis_type = "serendipity" + nb = gpython.basis.num_basis(basis_type, ndim, poly_order) + cells = [2, 1] + lower, upper = [0.0, 0.0], [2.0, 1.0] + + rng = np.random.default_rng(0) + coeffs = np.zeros((cells[0] * cells[1], nb)) + coeffs[0, :] = rng.normal(size=nb) + a = gpython.GkylArray.from_numpy(coeffs) + + grid = {"ndim": ndim, "lower": lower, "upper": upper, "cells": cells} + y0 = 0.3 + keep_dirs, cells_tar, out, btype, po_tar, cdim_tar, vdim_tar = ( + dg.modal.eval_at_coord_proj(grid, + basis_type, + ndim, + poly_order, + a, + eval_dirs=[1], + eval_coords=[y0])) + + assert keep_dirs == [0] + assert cells_tar == [2] + assert btype == "serendipity" + assert po_tar == 1 + + # cell 0 spans x in [0, 1], y in [0, 1]; center (0.5, 0.5), dx=dy=1. + x0 = 0.2 + zx = 2 * (x0 - 0.5) / 1.0 + zy = 2 * (y0 - 0.5) / 1.0 + donor_val = gpython.basis.eval_matrix(basis_type, ndim, poly_order, + np.array([[zx, zy]])) @ coeffs[0] + target_val = gpython.basis.eval_matrix(btype, 1, po_tar, np.array( + [[zx]])) @ out.view()[0] + np.testing.assert_allclose(target_val, donor_val, atol=1e-10) + + +@needs_gkeyll +def test_eval_at_coord_proj_full_reduction_uses_the_degenerate_1d_target(): + ndim, poly_order = 1, 1 + basis_type = "serendipity" + a = gpython.GkylArray.from_numpy(np.array([[1.0, 0.5]])) + grid = {"ndim": ndim, "lower": [0.0], "upper": [1.0], "cells": [1]} + + keep_dirs, cells_tar, out, btype, po_tar, cdim_tar, vdim_tar = ( + dg.modal.eval_at_coord_proj(grid, + basis_type, + ndim, + poly_order, + a, + eval_dirs=[0], + eval_coords=[0.5])) + + assert keep_dirs == [] + assert cells_tar == [1] + assert out.size == 1 + + +@needs_gkeyll +def test_eval_at_coord_proj_rejects_eval_dirs_out_of_range(): + a = gpython.GkylArray.alloc(gpython.basis.num_basis("serendipity", 2, 1), 1) + grid = {"ndim": 2, "lower": [0.0, 0.0], "upper": [1.0, 1.0], "cells": [1, 1]} + with pytest.raises(ValueError, match="out of range"): + dg.modal.eval_at_coord_proj(grid, + "serendipity", + 2, + 1, + a, + eval_dirs=[2], + eval_coords=[0.0]) + + +@needs_gkeyll +def test_eval_at_coord_proj_on_real_gkhybrid_data_matches_donor_reconstruction( +): + """Real 1x2v gkhybrid data: eliminating the vpar direction should produce + a target whose basis Gkeyll reports (possibly a different TYPE than the + donor's, since no gkhybrid convention has a mu-only velocity space), and + whose reconstructed value matches the donor's at the same physical point.""" + d = pg.load(GKHYB) + assert d.ctx["basis_type"] == "gkhybrid" + lo, up = d.bounds + y0 = 0.5 * (lo[1] + up[1]) # a vpar coordinate inside the domain + + out = d.eval_at_coord_proj([1], [y0]) + assert out.num_dims == 2 + + cells = np.asarray(d.ctx["cells"]) + dx = (up - lo) / cells + # Pick a point in cell (0, 0, 0): conf x and mu at their cell centers. + zx = 0.0 # cell center -> reference coordinate 0 + zy = 2.0 * (y0 - (lo[1] + 0.5 * dx[1])) / dx[1] + zmu = 0.0 + + donor_val = gpython.basis.eval_matrix("gkhybrid", 3, 1, + np.array([[zx, zy, zmu]])) @ np.asarray( + d.native.view())[0] + target_val = gpython.basis.eval_matrix( + out.ctx["basis_type"], 2, out.ctx["poly_order"], np.array( + [[zx, zmu]])) @ np.asarray(out.native.view())[0] + np.testing.assert_allclose(target_val, donor_val, atol=1e-8) + + +# =================================================== operations.eval_at_coord_proj +@needs_gkeyll +def test_ops_eval_at_coord_proj_rejects_numpy_backed_and_non_modal(): + interpolated = pg.load(GKHYB).interpolate() + with pytest.raises(ValueError, match="native modal data"): + interpolated.eval_at_coord_proj([1], [0.0]) + + nodal = pg.load(GKHYB).to_nodal() + with pytest.raises(ValueError, match="modal value_form"): + nodal.eval_at_coord_proj([1], [0.0]) + + +@needs_gkeyll +def test_ops_eval_at_coord_proj_rejects_missing_basis_metadata(): + a = pg.load(GKHYB) + del a.ctx["poly_order"] + with pytest.raises(ValueError, match="basis_type/poly_order"): + a.eval_at_coord_proj([1], [0.0]) + + +@needs_gkeyll +def test_ops_eval_at_coord_proj_tag_label_inplace(): + a = pg.load(GKHYB) + lo, up = a.bounds + y0 = 0.5 * (lo[1] + up[1]) + out = a.eval_at_coord_proj([1], [y0], tag="reduced", label="my label") + assert out.tag == "reduced" + assert out.label == "my label" + assert a.num_dims == 3 # original untouched + + b = pg.load(GKHYB) + mutated = b.eval_at_coord_proj([1], [y0], inplace=True) + assert mutated is b + assert b.num_dims == 2 + + +@needs_gkeyll +def test_ops_eval_at_coord_proj_can_eliminate_every_dimension(): + data = pg.load(GKHYB) + lower, upper = data.bounds + coordinates = 0.5 * (lower + upper) + + out = data.eval_at_coord_proj([0, 1, 2], coordinates) + + assert out.num_dims == 1 + assert out.ctx["cells"].tolist() == [1] + np.testing.assert_array_equal(out.grid[0], [0.0, 1.0]) diff --git a/tests/test_dg_map.py b/tests/test_dg_map.py new file mode 100644 index 00000000..ca0a9b54 --- /dev/null +++ b/tests/test_dg_map.py @@ -0,0 +1,219 @@ +"""Tests for ``postgkyl.dg.map`` -- grid mapping by evaluation at target points. + +See ``MAPPING.md`` for the design. Test fixtures build modal (or nodal) +coefficients for the mapping field synthetically with ``gpython.basis`` matrices +(no mapc2p file is required, per the layer instructions), by exactly +projecting a chosen physical-coordinate function onto the basis's own node +points, per cell -- this guarantees the coefficients exactly represent the +chosen function, so the expected result can be computed independently +(directly from the function), never from the code under test. + +Run: PYTHONPATH=src pytest tests/test_dg_map.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython, dg # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") +pytestmark = needs_gkeyll + + +def _project_1d(fn, lower, upper, cells, basis_type, poly_order): + """Exact per-cell modal coefficients of ``fn(z)`` for a 1-D basis. + + ``fn`` must be exactly representable in the basis on every cell (e.g. any + polynomial of degree <= poly_order) -- projecting through the node points + and back through the exact nodal<->modal change of basis reproduces it + exactly, independent of the mapping code under test. + """ + node_eta = gpython.basis.node_coords(basis_type, 1, poly_order)[:, 0] + n2m = gpython.basis.nodal_to_modal_matrix(basis_type, 1, poly_order) + dz = (upper - lower) / cells + centers = lower + (np.arange(cells) + 0.5) * dz + nodal_z = centers[:, None] + 0.5 * dz * node_eta[None, :] # (cells, nb) + nodal_vals = fn(nodal_z) + return nodal_vals @ n2m.T, nodal_vals # (modal, nodal) both (cells, nb) + + +def _project_2d(fn, lower, upper, cells, basis_type, poly_order): + """Exact per-cell modal coefficients of ``fn(z0, z1)`` for a 2-D basis.""" + node_eta = gpython.basis.node_coords(basis_type, 2, poly_order) # (nb, 2) + n2m = gpython.basis.nodal_to_modal_matrix(basis_type, 2, poly_order) + dz = [(upper[d] - lower[d]) / cells[d] for d in range(2)] + c0 = lower[0] + (np.arange(cells[0]) + 0.5) * dz[0] + c1 = lower[1] + (np.arange(cells[1]) + 0.5) * dz[1] + centers = np.stack(np.meshgrid(c0, c1, indexing="ij"), axis=-1) # (*cells,2) + # physical coordinates of every node, every cell: (*cells, nb, 2) + node_phys = ( + centers[:, :, None, :] + + 0.5 * np.array(dz)[None, None, None, :] * node_eta[None, None, :, :]) + nodal_vals = fn(node_phys[..., 0], node_phys[..., 1]) # (*cells, nb) + modal = np.einsum("ij,...j->...i", n2m, nodal_vals) + return modal, nodal_vals + + +# --------------------------------------------------------------------- 1-D +def test_eval_at_points_identity_map_1d_is_exact_to_machine_precision(): + lower, upper, cells = 0.0, 4.0, 4 + modal, _ = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + targets = np.linspace(lower, upper, 33) # finer than the mapping's own grid + got = dg.map.eval_at_points(modal, [lower], [upper], [cells], + targets[:, None], + basis_type="serendipity", + poly_order=1) + np.testing.assert_allclose(got, targets, atol=1e-12) + + +def test_map_grid_identity_1d_matches_target_axis(): + lower, upper, cells = -1.0, 3.0, 5 + modal, _ = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + target_axes = [np.linspace(lower, upper, 17)] + map_ctx = dict(lower=np.array([lower]), + upper=np.array([upper]), + cells=np.array([cells]), + basis_type="serendipity", + poly_order=1, + value_form="modal") + out = dg.map_grid(modal, map_ctx, target_axes) + assert len(out) == 1 + assert out[0].shape == target_axes[0].shape # m == 1 stays 1-D + np.testing.assert_allclose(out[0], target_axes[0], atol=1e-12) + + +def test_eval_at_points_in_basis_quadratic_is_exact_at_edges(): + """A single cell spanning the whole domain sidesteps cell-boundary + continuity questions entirely, isolating the eval_matrix/reshape math.""" + lower, upper, cells = -1.0, 3.0, 1 + fn = lambda z: 0.5 * z**2 - z + 1.0 # degree 2, in-basis for p2 + modal, _ = _project_1d(fn, lower, upper, cells, "serendipity", 2) + targets = np.array([lower, -0.3, 0.7, 2.1, upper]) # includes both edges + got = dg.map.eval_at_points(modal, [lower], [upper], [cells], + targets[:, None], + basis_type="serendipity", + poly_order=2) + np.testing.assert_allclose(got, fn(targets), atol=1e-12) + + +def test_eval_at_points_rejects_cells_mismatch(): + modal, _ = _project_1d(lambda z: z, 0.0, 1.0, 2, "serendipity", 1) + with pytest.raises(ValueError, match="does not match cells"): + dg.map.eval_at_points( + modal, + [0.0], + [1.0], + [3], # wrong cell count + np.array([[0.5]]), + basis_type="serendipity", + poly_order=1) + + +def test_eval_at_points_rejects_points_dim_mismatch(): + modal, _ = _project_1d(lambda z: z, 0.0, 1.0, 2, "serendipity", 1) + with pytest.raises(ValueError, match="expected 1"): + dg.map.eval_at_points( + modal, + [0.0], + [1.0], + [2], + np.array([[0.5, 0.5]]), # last axis length 2, expected 1 + basis_type="serendipity", + poly_order=1) + + +def test_eval_at_points_nodal_basis_path_matches_modal(): + """A nodal-basis mapping file (``nodal=True``) converts through the exact + nodal<->modal change of basis, then evaluates identically to the modal path.""" + lower, upper, cells = 0.0, 4.0, 4 + modal, nodal = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + targets = np.linspace(lower, upper, 11) + got_modal = dg.map.eval_at_points(modal, [lower], [upper], [cells], + targets[:, None], + basis_type="serendipity", + poly_order=1, + nodal=False) + got_nodal = dg.map.eval_at_points(nodal, [lower], [upper], [cells], + targets[:, None], + basis_type="serendipity", + poly_order=1, + nodal=True) + np.testing.assert_allclose(got_nodal, got_modal, atol=1e-12) + np.testing.assert_allclose(got_nodal, targets, atol=1e-12) + + +def test_map_grid_nodal_basis_map_file(): + lower, upper, cells = 0.0, 2.0, 2 + _, nodal = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + target_axes = [np.linspace(lower, upper, 9)] + map_ctx = dict(lower=np.array([lower]), + upper=np.array([upper]), + cells=np.array([cells]), + basis_type="serendipity", + poly_order=1, + value_form="nodal") + out = dg.map_grid(nodal, map_ctx, target_axes) + np.testing.assert_allclose(out[0], target_axes[0], atol=1e-12) + + +# --------------------------------------------------------------------- 2-D +def test_map_grid_identity_2d_curvilinear_matches_meshgrid(): + """Every physical coordinate is evaluated over all m dims, so the same + algorithm handles the non-separable (curvilinear) case.""" + lower, upper, cells = [0.0, 0.0], [2.0, 3.0], [2, 3] + m0, _ = _project_2d(lambda z0, z1: z0, lower, upper, cells, "serendipity", 1) + m1, _ = _project_2d(lambda z0, z1: z1, lower, upper, cells, "serendipity", 1) + map_coeffs = np.concatenate([m0, m1], axis=-1) + target_axes = [ + np.linspace(lower[0], upper[0], 5), + np.linspace(lower[1], upper[1], 7) + ] + map_ctx = dict(lower=np.array(lower), + upper=np.array(upper), + cells=np.array(cells), + basis_type="serendipity", + poly_order=1, + value_form="modal") + out = dg.map_grid(map_coeffs, map_ctx, target_axes) + + expected = np.meshgrid(*target_axes, indexing="ij") + assert len(out) == 2 + for d in range(2): + assert out[d].shape == (5, 7) # shape of the axes it replaces + np.testing.assert_allclose(out[d], expected[d], atol=1e-12) + + +def test_map_grid_2d_rotation_is_exact_non_separable(): + """A genuine rotation mixes both computational coordinates into each + physical one -- exercises the non-separable (curvilinear) evaluation.""" + lower, upper, cells = [-1.0, -1.0], [1.0, 1.0], [2, 2] + theta = 0.4 + cos_t, sin_t = np.cos(theta), np.sin(theta) + fn0 = lambda z0, z1: cos_t * z0 - sin_t * z1 + fn1 = lambda z0, z1: sin_t * z0 + cos_t * z1 + m0, _ = _project_2d(fn0, lower, upper, cells, "serendipity", 1) + m1, _ = _project_2d(fn1, lower, upper, cells, "serendipity", 1) + map_coeffs = np.concatenate([m0, m1], axis=-1) + target_axes = [ + np.linspace(lower[0], upper[0], 6), + np.linspace(lower[1], upper[1], 4) + ] + map_ctx = dict(lower=np.array(lower), + upper=np.array(upper), + cells=np.array(cells), + basis_type="serendipity", + poly_order=1, + value_form="modal") + out = dg.map_grid(map_coeffs, map_ctx, target_axes) + + z0, z1 = np.meshgrid(*target_axes, indexing="ij") + np.testing.assert_allclose(out[0], fn0(z0, z1), atol=1e-12) + np.testing.assert_allclose(out[1], fn1(z0, z1), atol=1e-12) diff --git a/tests/test_dg_rep.py b/tests/test_dg_rep.py new file mode 100644 index 00000000..c03e5633 --- /dev/null +++ b/tests/test_dg_rep.py @@ -0,0 +1,106 @@ +"""Tests for ``postgkyl.dg.rep`` -- modal · nodal · quad representation changes. + +This is the module's dedicated home post-relocation (``gpython/rep.py`` -> +``dg/rep.py``, layer 03-dg job 1); defensive/edge-case branches for the same +module are also exercised from ``tests/test_coverage_leaf.py`` (a shared +leaf/engine coverage file predating this move). See ``CLAUDE.md``'s "Engine +layers" section for why representation changes live in ``dg``, not ``gpython``. + +Run: PYTHONPATH=src pytest tests/test_dg_rep.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython, dg # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") +pytestmark = needs_gkeyll + + +def _linear_field(basis_type, ndim, poly_order, cells, nfields=1): + """An exactly-representable modal field: coefficient 0 (mean) = cell index, + everything else zero -- lets every conversion be checked against a value + known independently of the code under test.""" + nb = gpython.basis.num_basis(basis_type, ndim, poly_order) + ncells = int(np.prod(cells)) + vals = np.zeros((ncells, nfields * nb)) + for f in range(nfields): + vals[:, f * nb] = np.arange(ncells) + f # only the mean coefficient + return gpython.GkylArray.from_numpy(vals), nb + + +def test_modal_to_nodal_to_modal_round_trips_exactly(): + arr, nb = _linear_field("serendipity", 2, 1, [3, 3]) + nodal = dg.rep.modal_to_nodal("serendipity", 2, 1, arr) + back = dg.rep.nodal_to_modal("serendipity", 2, 1, nodal) + np.testing.assert_allclose(back.view(), arr.view(), atol=1e-13) + + +def test_modal_to_nodal_matches_a_directly_evaluated_constant_field(): + """A pure mean-coefficient field is a constant per cell -- every nodal + value must equal that constant, checked against the analytic normalized + constant basis function b0 = 2^(-ndim/2), independent of the shim.""" + ndim, poly_order = 1, 1 + arr, nb = _linear_field("serendipity", ndim, poly_order, [4]) + nodal = dg.rep.modal_to_nodal("serendipity", ndim, poly_order, arr) + b0 = 2.0**(-ndim / 2.0) + expected = (np.arange(4) * b0)[:, None] * np.ones(nb) + np.testing.assert_allclose(nodal.view(), expected, atol=1e-13) + + +@pytest.mark.parametrize("num_quad", [2, 3]) +def test_quad_round_trip_exact_for_in_basis_field(num_quad): + """modal -> quad -> modal is exact whenever num_quad >= poly_order + 1.""" + arr, nb = _linear_field("serendipity", 1, 1, [5], nfields=2) + quad = dg.rep.modal_to_quad("serendipity", 1, 1, arr, num_quad) + back = dg.rep.quad_to_modal("serendipity", 1, 1, quad, num_quad) + np.testing.assert_allclose(back.view(), arr.view(), atol=1e-12) + + +def test_wrap_round_trips_values_unchanged(): + values = np.arange(12.0).reshape(3, 4) + wrapped = dg.rep.wrap(values) + np.testing.assert_array_equal(wrapped.view(), values) + + +def test_apply_pointwise_sqrt_matches_numpy_after_interpolate(): + """fn applied via quadrature matches applying fn directly to the exact + (interpolated) values, for an in-basis-representable nonnegative field.""" + ndim, poly_order = 1, 1 + arr, nb = _linear_field("serendipity", ndim, poly_order, [4]) + arr = gpython.kernels.shiftc(arr, 5.0, 0) # keep the field positive for sqrt + out = dg.rep.apply_pointwise(ndim=ndim, + poly_order=poly_order, + basis_type="serendipity", + arr=arr, + fn=np.sqrt, + num_quad=poly_order + 1) + grid, direct_vals = dg.interpolate(arr.view(), [np.linspace(0, 4, 5)], + poly_order=poly_order, + basis_type="serendipity") + grid, sqrt_of_applied = dg.interpolate(out.view(), [np.linspace(0, 4, 5)], + poly_order=poly_order, + basis_type="serendipity") + np.testing.assert_allclose(sqrt_of_applied, np.sqrt(direct_vals), atol=1e-10) + + +def test_materialize_nodal_matches_modal_to_nodal_values(): + ndim, poly_order = 1, 1 + arr, nb = _linear_field("serendipity", ndim, poly_order, [3]) + nodal = dg.rep.modal_to_nodal("serendipity", ndim, poly_order, arr) + grid = [np.linspace(0.0, 3.0, 4)] + edges, out = dg.rep.materialize("serendipity", ndim, poly_order, nodal, grid, + "nodal") + assert out.shape == (2 * 3, 1) # 2 tensor nodes/cell * 3 cells, 1 field + np.testing.assert_allclose(np.sort(out[:, 0].reshape(3, 2), axis=1), + np.sort(nodal.view(), axis=1), + atol=1e-13) diff --git a/tests/test_diagnostics_discovery.py b/tests/test_diagnostics_discovery.py new file mode 100644 index 00000000..5dc6f8ec --- /dev/null +++ b/tests/test_diagnostics_discovery.py @@ -0,0 +1,81 @@ +"""Tests for postgkyl.diagnostics.discovery -- the equation-blind +output-stem/frame discovery shared by every equation loader. + +No dedicated ``find_output_stems``/``.outputs()`` tests exist in +``tests_bak`` (``tests_bak/test_loader.py`` tests the ``pg.load`` +callable/namespace instead -- see ``test_diagnostics_gk_load.py``'s +``TestResolveFrames`` for the pieces of that file that do belong to this +layer), so this is a fresh corpus targeting ``find_output_stems`` and the new +``available_frames`` helper directly. +""" + +from __future__ import annotations + +from postgkyl.diagnostics import discovery + + +def _touch(tmp_path, *names): + for name in names: + (tmp_path / name).touch() + + +class TestFindOutputStems: + + def test_single_extension_single_stem(self, tmp_path): + _touch(tmp_path, "elc_M0_0.gkyl", "elc_M0_1.gkyl", "elc_M0_2.gkyl") + out = discovery.find_output_stems("gkyl", str(tmp_path)) + assert out == {"gkyl": ["elc_M0"]} + + def test_multiple_stems_sorted(self, tmp_path): + _touch(tmp_path, "ion_M0_0.gkyl", "elc_M0_0.gkyl", "field_0.gkyl") + out = discovery.find_output_stems("gkyl", str(tmp_path)) + assert out["gkyl"] == ["elc_M0", "field", "ion_M0"] + + def test_multiple_extensions(self, tmp_path): + _touch(tmp_path, "elc_M0_0.gkyl", "elc_M0_0.h5") + out = discovery.find_output_stems("h5,gkyl", str(tmp_path)) + assert out == {"h5": ["elc_M0"], "gkyl": ["elc_M0"]} + + def test_strips_restart_suffix(self, tmp_path): + _touch(tmp_path, "elc_M0_0_restart.gkyl") + out = discovery.find_output_stems("gkyl", str(tmp_path)) + assert out["gkyl"] == ["elc_M0"] + + def test_no_frame_number_kept_as_is(self, tmp_path): + _touch(tmp_path, "geo_int_jacobtot_inv.gkyl") + out = discovery.find_output_stems("gkyl", str(tmp_path)) + assert out["gkyl"] == ["geo_int_jacobtot_inv"] + + def test_empty_directory(self, tmp_path): + out = discovery.find_output_stems("gkyl", str(tmp_path)) + assert out == {"gkyl": []} + + def test_default_extensions_and_path(self, tmp_path, monkeypatch): + _touch(tmp_path, "a_0.gkyl") + monkeypatch.chdir(tmp_path) + out = discovery.find_output_stems() + assert out == {"gkyl": ["a"]} + + +class TestAvailableFrames: + + def test_discovers_all_frames(self, tmp_path): + stem = str(tmp_path / "sim-ion_M0_") + _touch(tmp_path, "sim-ion_M0_0.gkyl", "sim-ion_M0_1.gkyl", + "sim-ion_M0_5.gkyl") + assert discovery.available_frames(stem) == {0, 1, 5} + + def test_restricted_to_candidate_frames(self, tmp_path): + stem = str(tmp_path / "sim-ion_M0_") + _touch(tmp_path, "sim-ion_M0_0.gkyl", "sim-ion_M0_1.gkyl", + "sim-ion_M0_5.gkyl") + assert discovery.available_frames(stem, frames=[0, 5, 99]) == {0, 5} + + def test_no_matching_files(self, tmp_path): + stem = str(tmp_path / "sim-ion_M0_") + assert discovery.available_frames(stem) == set() + + def test_non_numeric_suffix_ignored(self, tmp_path): + stem = str(tmp_path / "sim-ion_M0_") + _touch(tmp_path, "sim-ion_M0_0.gkyl", "sim-ion_M0_restart.gkyl") + assert discovery.available_frames(stem) == {0} diff --git a/tests/test_diagnostics_five_moment.py b/tests/test_diagnostics_five_moment.py new file mode 100644 index 00000000..b166ce50 --- /dev/null +++ b/tests/test_diagnostics_five_moment.py @@ -0,0 +1,295 @@ +"""Tests for postgkyl.diagnostics.mom.five_moment -- the 5-/10-moment primitive +variable family (density, velocity, pressure, temperature, sound, Mach), +folding the array-math analytic tests (formerly tests_models_five_moment.py) +with the verb-level guard/inplace/tag/label/VARIABLES tests (formerly part of +tests_ops_moments.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.mom import five_moment as fm +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +_G1D = [np.array([0.0, 1.0])] + +# 5-moment Euler fluid: [rho, rho*vx, rho*vy, rho*vz, E] +_RHO = 1.0 +_VX, _VY, _VZ = 0.5, 0.25, 0.1 +_P_THERMAL = 0.6 +_GAMMA = 5.0 / 3.0 +_E_5 = _P_THERMAL / (_GAMMA - 1) + 0.5 * _RHO * (_VX**2 + _VY**2 + _VZ**2) +_MOM5 = np.array([[_RHO, _RHO * _VX, _RHO * _VY, _RHO * _VZ, _E_5]]) + +# 10-moment fluid: [rho, mx, my, mz, Pxx, Pxy, Pxz, Pyy, Pyz, Pzz] +_P_T = 0.4 +_Pxx = _P_T + _RHO * _VX**2 +_Pxy = 0.0 + _RHO * _VX * _VY +_Pxz = 0.0 + _RHO * _VX * _VZ +_Pyy = _P_T + _RHO * _VY**2 +_Pyz = 0.0 + _RHO * _VY * _VZ +_Pzz = _P_T + _RHO * _VZ**2 +_MOM10 = np.array([[ + _RHO, _RHO * _VX, _RHO * _VY, _RHO * _VZ, _Pxx, _Pxy, _Pxz, _Pyy, _Pyz, _Pzz +]]) + + +class TestDensity: + + def test_value(self): + d = _make(_G1D, _MOM5) + out = fm.density(d) + np.testing.assert_allclose(out.values[0, 0], _RHO) + + def test_output_shape_has_trailing_dim(self): + d = _make(_G1D, _MOM5) + out = fm.density(d) + assert out.values.ndim == _MOM5.ndim + assert out.values.shape[-1] == 1 + + def test_multi_cell(self): + grid = [np.linspace(0.0, 1.0, 4)] + values = np.hstack([np.array([[1.0], [2.0], [3.0]]), np.zeros((3, 4))]) + d = _make(grid, values) + out = fm.density(d) + np.testing.assert_allclose(out.values[:, 0], [1.0, 2.0, 3.0]) + + def test_inplace_mutates(self): + d = _make(_G1D, _MOM5) + out = fm.density(d, inplace=True) + assert out is d + + def test_tag_and_label(self): + d = _make(_G1D, _MOM5) + out = fm.density(d, tag="rho", label="lbl") + assert out.get_tag() == "rho" + assert out.get_label() == "lbl" + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.density(d) + + +class TestVelocityComponents: + + def test_xvel(self): + d = _make(_G1D, _MOM5) + out = fm.xvel(d) + np.testing.assert_allclose(out.values[0, 0], _VX) + + def test_yvel(self): + d = _make(_G1D, _MOM5) + out = fm.yvel(d) + np.testing.assert_allclose(out.values[0, 0], _VY) + + def test_zvel(self): + d = _make(_G1D, _MOM5) + out = fm.zvel(d) + np.testing.assert_allclose(out.values[0, 0], _VZ) + + def test_vel_three_components(self): + d = _make(_G1D, _MOM5) + out = fm.vel(d) + assert out.values.shape[-1] == 3 + np.testing.assert_allclose(out.values[0, 0], _VX) + np.testing.assert_allclose(out.values[0, 1], _VY) + np.testing.assert_allclose(out.values[0, 2], _VZ) + + def test_fabricated_maxwellian_recovers_bulk_velocity(self): + # density=1, momentum=(2, 0, 0), energy=10: analytic case from the + # legacy TestMomentFluent euler() fixture -- vx should recover 2.0. + d = _make([np.array([0.0, 1.0])], np.array([[1.0, 2.0, 0.0, 0.0, 10.0]])) + rho_out = fm.density(d) + vx_out = fm.xvel(d) + np.testing.assert_allclose(rho_out.values.flat[0], 1.0) + np.testing.assert_allclose(vx_out.values.flat[0], 2.0) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.xvel(d) + + +class TestPressureScalar: + + def test_5mom_auto_detect(self): + d = _make(_G1D, _MOM5) + out = fm.pressure(d) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL, rtol=1e-10) + + def test_5mom_explicit(self): + d = _make(_G1D, _MOM5) + out = fm.pressure(d, num_moms=5) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL, rtol=1e-10) + + def test_10mom_auto_detect(self): + d = _make(_G1D, _MOM10) + out = fm.pressure(d) + np.testing.assert_allclose(out.values[0, 0], _P_T, rtol=1e-10) + + def test_10mom_explicit(self): + d = _make(_G1D, _MOM10) + out = fm.pressure(d, num_moms=10) + np.testing.assert_allclose(out.values[0, 0], _P_T, rtol=1e-10) + + def test_wrong_num_comps_raises(self): + d = _make(_G1D, np.array([[1.0, 2.0, 3.0]])) + with pytest.raises(ValueError, match="num_moms"): + fm.pressure(d) + + def test_multi_cell(self): + grid = [np.linspace(0.0, 1.0, 3)] + values = np.concatenate([_MOM5, _MOM5 * 2.0], axis=0) + d = _make(grid, values) + out = fm.pressure(d, num_moms=5) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL, rtol=1e-9) + np.testing.assert_allclose(out.values[1, 0], 2.0 * _P_THERMAL, rtol=1e-9) + + def test_gas_gamma_is_forwarded(self): + d = _make(_G1D, _MOM5) + out = fm.pressure(d, gas_gamma=1.4) + _, expected = fm._get_p(d.grid, d.values, gas_gamma=1.4, num_moms=5) + np.testing.assert_allclose(out.values, expected) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.pressure(d) + + +class TestKineticEnergy: + + def test_5mom(self): + d = _make(_G1D, _MOM5) + out = fm.ke(d) + expected = 0.5 * _RHO * (_VX**2 + _VY**2 + _VZ**2) + np.testing.assert_allclose(out.values[0, 0], expected, rtol=1e-10) + + def test_10mom(self): + d = _make(_G1D, _MOM10) + out = fm.ke(d, num_moms=10) + expected = 0.5 * _RHO * (_VX**2 + _VY**2 + _VZ**2) + np.testing.assert_allclose(out.values[0, 0], expected, rtol=1e-10) + + def test_wrong_num_comps_raises(self): + d = _make(_G1D, np.array([[1.0, 2.0, 3.0]])) + with pytest.raises(ValueError): + fm.ke(d) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.ke(d) + + +class TestTempSoundMach: + + def test_temp_5mom(self): + d = _make(_G1D, _MOM5) + out = fm.temp(d) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL / _RHO, rtol=1e-10) + + def test_temp_10mom(self): + d = _make(_G1D, _MOM10) + out = fm.temp(d, num_moms=10) + np.testing.assert_allclose(out.values[0, 0], _P_T / _RHO, rtol=1e-10) + + def test_sound_speed(self): + d = _make(_G1D, _MOM5) + out = fm.sound(d) + expected = np.sqrt(_GAMMA * _P_THERMAL / _RHO) + np.testing.assert_allclose(out.values[0, 0], expected, rtol=1e-10) + + def test_mach(self): + d = _make(_G1D, _MOM5) + out = fm.mach(d) + v = np.sqrt(_VX**2 + _VY**2 + _VZ**2) + cs = np.sqrt(_GAMMA * _P_THERMAL / _RHO) + np.testing.assert_allclose(out.values[0, 0], v / cs, rtol=1e-10) + + def test_grid_is_passed_through_unchanged(self): + d = _make(_G1D, _MOM5) + out = fm.mach(d) + np.testing.assert_allclose(out.grid[0], _G1D[0]) + + @needs_gkeyll + def test_temp_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.temp(d) + + @needs_gkeyll + def test_sound_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.sound(d) + + @needs_gkeyll + def test_mach_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.mach(d) + + +class TestVelocityVerb: + + def test_divides_momentum_by_density(self): + density = _make([np.array([0.0, 1.0, 2.0])], np.array([[1.0], [2.0]])) + momentum = _make([np.array([0.0, 1.0, 2.0])], + np.array([[3.0, 6.0], [4.0, 8.0]])) + out = fm.velocity(density, momentum) + np.testing.assert_allclose(out.values, [[3.0, 6.0], [2.0, 4.0]]) + + def test_inplace_mutates_density(self): + density = _make([np.array([0.0, 1.0])], np.array([[1.0]])) + momentum = _make([np.array([0.0, 1.0])], np.array([[2.0]])) + out = fm.velocity(density, momentum, inplace=True) + assert out is density + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + field = _make([np.array([0.0, 1.0])], np.array([[1.0]])) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fm.velocity(d, field) + + +class TestVariables: + + @pytest.mark.parametrize("name", [ + "density", "xvel", "yvel", "zvel", "vel", "pressure", "ke", "temp", + "sound", "mach" + ]) + def test_variables_table_matches_public_function(self, name): + assert fm.VARIABLES[name] is getattr(fm, name) + + def test_variables_table_has_exactly_the_old_euler_vocabulary(self): + assert set(fm.VARIABLES) == { + "density", "xvel", "yvel", "zvel", "vel", "pressure", "ke", "temp", + "sound", "mach" + } diff --git a/tests/test_diagnostics_gk_load.py b/tests/test_diagnostics_gk_load.py new file mode 100644 index 00000000..5cbad1ad --- /dev/null +++ b/tests/test_diagnostics_gk_load.py @@ -0,0 +1,1147 @@ +"""Tests for the gyrokinetic loader stack: +``postgkyl.diagnostics.gk.{distf,quantity,quantities,registry, +load_quantity}``. + +Ported/extended from ``tests_bak/test_gk_load_quantity.py`` (the registry +smoke test, using the same "synthetic constant DG field + monkeypatched +``GData``" technique) and the ``TestResolveFrames``/``TestLoadGkDistf`` +classes of ``tests_bak/test_loader.py`` (``pg.load.gk_distf``'s dispatch +tests do not port: this architecture has no ``pg.load`` namespace object -- +``load_distf``/``resolve_frames`` are plain free functions, tested +directly). Real end-to-end coverage uses the ``rt_gk_tcv_iwl*`` fixtures +staged in ``tests/test_data`` for this layer. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from postgkyl import gpython +from postgkyl.gdata import GData, GDataGroup +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.diagnostics.gk import distf, quantities as ff, quantity as qmod, utils +from postgkyl.diagnostics.gk.load_quantity import (available_quantities, + load_quantity) +from postgkyl.diagnostics.gk.registry import gk_quant_registry + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GK_NAME = "rt_gk_tcv_iwl_1x2v_p1" +HMOM_NAME = "rt_gk_tcv_iwl_adapt_source_1x2v_p1" + + +def _field(values, grid=None, **ctx): + """A pre-interpolated (field-domain) dataset for unit-testing the + ``fetch_*`` combinators without needing the compiled shim.""" + d = GDataState(ctx=dict(ctx, interpolated=True)) + values = np.asarray(values, dtype=np.float64) + if grid is None: + grid = [ + np.arange(values.shape[ax] + 1, dtype=np.float64) + for ax in range(values.ndim - 1) + ] + d.push(grid, values) + return d + + +class TestResolveFrames: + """Ported from tests_bak/test_loader.py's TestResolveFrames.""" + + def test_single_int(self): + assert distf.resolve_frames(5, name="n", species="ion") == [5] + + def test_list(self): + assert distf.resolve_frames([1, 2, 3], name="n", species="ion") == [1, 2, 3] + + def test_csv_string(self): + assert distf.resolve_frames("0,2,4", name="n", species="ion") == [0, 2, 4] + + def test_single_element_list(self): + assert distf.resolve_frames([7], name="n", species="ion") == [7] + + def test_range_discovers_files(self, tmp_path, monkeypatch): + for f in (0, 1, 2, 3): + (tmp_path / f"sim-ion_{f}.gkyl").touch() + monkeypatch.chdir(tmp_path) + assert distf.resolve_frames("1:3", name="sim", species="ion") == [1, 2] + assert distf.resolve_frames(":", name="sim", species="ion") == [0, 1, 2, 3] + assert distf.resolve_frames("0:4:2", name="sim", species="ion") == [0, 2] + + def test_numeric_string(self): + assert distf.resolve_frames("7", name="n", species="ion") == [7] + + def test_range_without_matching_files_has_a_clear_error( + self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + with pytest.raises(ValueError, match="No distribution frames found"): + distf.resolve_frames(":", name="sim", species="ion") + + def test_range_requires_a_positive_step(self, tmp_path, monkeypatch): + (tmp_path / "sim-ion_0.gkyl").touch() + monkeypatch.chdir(tmp_path) + with pytest.raises(ValueError, match="positive integer"): + distf.resolve_frames("::0", name="sim", species="ion") + + +class TestLoadGkDistfFrames: + """The public loader resolves frame syntax around the per-frame core.""" + + def _stub(self, monkeypatch): + calls = [] + + def fake_load_distf_frame(*, frame, tag, **kwargs): + calls.append(frame) + data = GData(tag=tag, ctx={"frame": frame}) + data.push([np.array([0.0, 1.0])], np.array([[float(frame)]])) + return data + + monkeypatch.setattr(distf, "_load_distf_frame", fake_load_distf_frame) + return calls + + def test_integer_frame_returns_one_dataset(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 3) + assert isinstance(out, GData) + assert calls == [3] + + def test_csv_frames_return_a_labelled_fluent_group(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", "0,2,4") + assert isinstance(out, GDataGroup) + assert calls == [0, 2, 4] + assert [data.label for data in out] == ["0", "2", "4"] + + def test_single_element_list_still_returns_a_group(self, monkeypatch): + self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", [7]) + assert isinstance(out, GDataGroup) + assert len(out) == 1 + + def test_range_loads_only_discovered_frames(self, tmp_path, monkeypatch): + calls = self._stub(monkeypatch) + for frame in (0, 2, 5): + (tmp_path / f"sim-ion_fdot_{frame}.gkyl").touch() + monkeypatch.chdir(tmp_path) + out = distf.load_distf("sim", "ion", "0:5", suffix="fdot") + assert isinstance(out, GDataGroup) + assert calls == [0, 2] + + +class TestLoadGkDistfKeywordOnly: + """``load_distf``'s options must be keyword-only (PYTHON_PRINCIPLES #7 / + doctrine IV) so a caller can never silently swap two boolean flags by + passing them positionally.""" + + def test_tag_cannot_be_passed_positionally(self): + with pytest.raises(TypeError): + distf.load_distf("sim", "ion", 0, "f") + + +@needs_gkeyll +class TestLoadGkDistfReal: + """End-to-end against the staged rt_gk_tcv_iwl_1x2v_p1 fixtures. + + ``mapc2p_vel``/``jacobvel`` in the fixture set carry no DG (basis_type/ + poly_order) metadata, so the coordinate-mapping options (``use_c2p_vel`` + etc., which need ``operations.map`` to read that metadata off the mapping file) + cannot be exercised against these particular files; only the default + (no-mapping) path is covered here. + """ + + def test_shape_and_grid(self): + out = distf.load_distf(name=os.path.join(DATA, GK_NAME), + species="elc", + frame=250, + jacobtot_inv_file=os.path.join( + DATA, f"{GK_NAME}-geo_int_jacobtot_inv.gkyl")) + assert out.num_dims == 3 + assert out.num_comps == 1 + assert out.values.shape[:3] == tuple(int(c) for c in out.num_cells) + assert np.all(np.isfinite(out.values)) + + def test_missing_jacobtot_inv_file_raises(self): + with pytest.raises(Exception): + distf.load_distf(name=os.path.join(DATA, GK_NAME), + species="elc", + frame=250, + jacobtot_inv_file=os.path.join(DATA, + "does_not_exist.gkyl")) + + +class _FakeDistfData(GData): + """A ``GData`` whose ``.interpolate()`` is a stubbed no-op (real + interpolation needs the compiled Gkeyll shim), keeping the real computing + operators (``*``/``/``) so ``load_distf``'s weak-multiply-then-divide + step still runs (as a plain NumPy op, since these fakes are never + gkyl-native) -- letting ``load_distf``'s coordinate-map branches + (``use_c2p_vel``/``use_mc2nu``/``use_mapc2p``) be exercised without real + mapc2p_vel/mc2nu/mapc2p DG fixtures (the staged rt_gk_tcv_iwl* files carry + no such metadata -- see TestLoadGkDistfReal).""" + + def interpolate(self, + *, + basis=None, + p=None, + num_interp=None, + inplace=False, + tag=None, + label=None): + return self + + +class TestLoadGkDistfCoordinateMaps: + """Unit tests of ``load_distf``'s ``use_c2p_vel``/``use_mc2nu``/ + ``use_mapc2p`` branches, stubbed through ``distf.load``/``operations.map`` + since the compiled-Gkeyll fixtures have no mapping-file metadata to + exercise them against.""" + + def _stub(self, monkeypatch): + grid = [np.linspace(0.0, 1.0, 5)] + values = np.ones((4, 1)) + registry = { + "sim-ion_0.gkyl": (grid, values), + "sim-ion_jacobvel.gkyl": (grid, values), + "sim-geo_int_jacobtot_inv.gkyl": (grid, values), + "sim-ion_mapc2p_vel.gkyl": (grid, values), + } + + def fake_load(file_name="", + *, + tag="default", + label="", + ctx=None, + value_form=None, + **read_kwargs): + d = _FakeDistfData(tag=tag, label=label, ctx=ctx) + if file_name: + d.push(*registry[file_name]) + d._file_name = file_name + return d + + monkeypatch.setattr(distf, "load", fake_load) + calls = [] + + def fake_map(data, mapping, *, space, basis_type=None, poly_order=None): + # mapc2p_vel is pre-loaded (to attach basis_type/poly_order overrides) + # before it reaches operations.map, so `mapping` arrives as a + # dataset there, not a filename -- unlike the conf-space maps + # (mc2nu/mapc2p), which are still passed through as bare paths. + recorded = mapping._file_name if hasattr(mapping, + "_file_name") else mapping + calls.append((recorded, space)) + return data + + monkeypatch.setattr(distf.operations, "map", fake_map) + return calls + + def test_use_c2p_vel(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 0, use_c2p_vel=True) + assert calls == [("sim-ion_mapc2p_vel.gkyl", "vel")] + assert out.ctx["grid_type"] == "c2p_vel" + + def test_use_mc2nu(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 0, use_mc2nu=True) + assert calls == [("sim-geo_corn_mc2nu_pos_deflated.gkyl", "conf")] + assert out.ctx["grid_type"] == "mc2nu" + + def test_use_mapc2p(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 0, use_mapc2p=True) + assert calls == [("sim-geo_corn_mapc2p_deflated.gkyl", "conf")] + assert out.ctx["grid_type"] == "mapc2p" + + def test_use_mc2nu_takes_precedence_over_mapc2p(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 0, use_mc2nu=True, use_mapc2p=True) + assert calls == [("sim-geo_corn_mc2nu_pos_deflated.gkyl", "conf")] + assert out.ctx["grid_type"] == "mc2nu" + + def test_use_c2p_vel_and_mapc2p_both_applied(self, monkeypatch): + calls = self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 0, use_c2p_vel=True, use_mapc2p=True) + assert calls == [("sim-ion_mapc2p_vel.gkyl", "vel"), + ("sim-geo_corn_mapc2p_deflated.gkyl", "conf")] + assert out.ctx["grid_type"] == "c2p_vel + mapc2p" + + def test_no_grid_type_key_when_no_maps_requested(self, monkeypatch): + self._stub(monkeypatch) + out = distf.load_distf("sim", "ion", 0) + assert "grid_type" not in out.ctx + + +class TestFetchCombinators: + """Unit tests of the generic component-extraction/combinator factories -- + pure field-domain math, no compiled shim needed.""" + + def test_component_extraction(self): + d = _field(np.array([[1.0, 2.0, 3.0]] * 3)) + out = ff._component(d, 1) + np.testing.assert_allclose(out.values[..., 0], 2.0) + + def test_component_all(self): + d = _field(np.array([[1.0, 2.0, 3.0]] * 3)) + out = ff._component(d, None) + assert out.values.shape[-1] == 3 + + def test_binop_add(self): + a = _field(np.array([[1.0, 10.0]] * 2)) + fetch = ff._make_fetch_binop(0, 0, 0, 1, lambda x, y: x + y) + out = fetch([a]) + np.testing.assert_allclose(out.values[..., 0], 11.0) + + def test_fetch_s1c0_div_s0c0(self): + m0 = _field(np.full((3, 1), 2.0)) + m1 = _field(np.full((3, 1), 6.0)) + out = ff.fetch_s1c0_div_s0c0([m0, m1]) + np.testing.assert_allclose(out.values, 3.0) + + +class TestFetchPhysics: + """Analytic checks of the derived-quantity formulas, using hand-built + field-domain fixtures (mass/charge as ctx or via the ``**extra`` fallback, + matching ``_get_ctx_val``'s contract).""" + + def test_M1_from_H(self): + hmom = _field(np.full((3, 2), 1.0), mass=2.0) + hmom.values[..., 0] = 4.0 + hmom.values[..., 1] = 3.0 + out = ff.fetch_M1_from_H([hmom]) + np.testing.assert_allclose(out.values[..., 0], 4.0 * 3.0 / 2.0) + + def test_Tpar_from_BiMax(self): + bimax = _field(np.zeros((2, 4)), mass=3.0) + bimax.values[..., 2] = 5.0 + out = ff.fetch_Tpar_from_BiMax([bimax]) + np.testing.assert_allclose(out.values[..., 0], 15.0) + + def test_Tpar_from_M0_M1_M2par(self): + m0 = _field(np.full((2, 1), 2.0), mass=4.0) + m1 = _field(np.full((2, 1), 6.0)) + m2par = _field(np.full((2, 1), 10.0)) + out = ff.fetch_Tpar_from_M0_M1_M2par([m0, m1, m2par]) + # Tpar = mass*(M2par - M1**2/M0)/M0 = 4*(10 - 36/2)/2 = 4*(-8)/2 = -16 + np.testing.assert_allclose(out.values[..., 0], -16.0) + + def test_temp_from_Tpar_Tperp(self): + Tpar = _field(np.full((2, 1), 3.0)) + Tperp = _field(np.full((2, 1), 6.0)) + out = ff.fetch_temp_from_Tpar_Tperp([Tpar, Tperp]) + np.testing.assert_allclose(out.values[..., 0], (3.0 + 2 * 6.0) / 3.0) + + def test_press_p(self): + m0 = _field(np.full((2, 1), 2.0)) + Tp = _field(np.full((2, 1), 5.0)) + out = ff.fetch_press_p([m0, Tp]) + np.testing.assert_allclose(out.values[..., 0], 10.0) + + def test_beta_from_bmag_press(self): + from scipy import constants + bmag = _field(np.full((2, 1), 2.0)) + press = _field(np.full((2, 1), 5.0)) + out = ff.fetch_beta_from_bmag_press([bmag, press]) + np.testing.assert_allclose(out.values[..., 0], + 2.0 * constants.mu_0 * 5.0 / 4.0) + + def test_missing_ctx_key_raises(self): + m0 = _field(np.full((2, 1), 2.0)) + with pytest.raises(KeyError): + ff.fetch_M1_from_H([m0]) + + def test_missing_ctx_key_uses_extra(self): + hmom = _field(np.full((2, 2), 1.0)) + hmom.values[..., 0] = 4.0 + hmom.values[..., 1] = 3.0 + out = ff.fetch_M1_from_H([hmom], mass=2.0) + np.testing.assert_allclose(out.values[..., 0], 6.0) + + def test_Tperp_from_M0_M2perp(self): + m0 = _field(np.full((2, 1), 2.0), mass=3.0) + m2perp = _field(np.full((2, 1), 8.0)) + out = ff.fetch_Tperp_from_M0_M2perp([m0, m2perp]) + # Tperp = 0.5*mass*(M2perp/M0) = 0.5*3*(8/2) = 6 + np.testing.assert_allclose(out.values[..., 0], 6.0) + + def test_temp_from_Max(self): + maxmom = _field(np.zeros((2, 3)), mass=2.0) + maxmom.values[..., 2] = 5.0 + out = ff.fetch_temp_from_Max([maxmom]) + np.testing.assert_allclose(out.values[..., 0], 10.0) + + def test_press_from_Max(self): + maxmom = _field(np.zeros((2, 3)), mass=2.0) + maxmom.values[..., 0] = 3.0 + maxmom.values[..., 2] = 5.0 + out = ff.fetch_press_from_Max([maxmom]) + np.testing.assert_allclose(out.values[..., 0], 2.0 * 3.0 * 5.0) + + def test_press_from_BiMax(self): + bimax = _field(np.zeros((2, 4)), mass=2.0) + bimax.values[..., 0] = 3.0 # M0 + bimax.values[..., 2] = 4.0 # Tpar (pre-mass) + bimax.values[..., 3] = 5.0 # Tperp (pre-mass) + out = ff.fetch_press_from_BiMax([bimax]) + # press = M0 * mass*(Tpar + 2*Tperp)/3 = 3 * 2*(4 + 10)/3 = 3*28/3 = 28 + np.testing.assert_allclose(out.values[..., 0], 28.0) + + +class TestGetCtxVal: + """Resolution of species attributes: an explicit '--extra' override wins + over the file's own context, which wins over raising; an '--extra' value + may be a single scalar (every species) or one entry per species.""" + + def test_extra_overrides_the_context(self): + d = _field(np.full((2, 1), 1.0), mass=5.0) + assert ff._get_ctx_val(d, "mass", mass=999.0) == 999.0 + + def test_context_is_used_when_extra_does_not_carry_the_key(self): + d = _field(np.full((2, 1), 1.0), mass=5.0) + assert ff._get_ctx_val(d, "mass") == 5.0 + assert ff._get_ctx_val(d, "mass", charge=1.0) == 5.0 + + def test_scalar_extra_applies_to_every_species(self): + d = _field(np.full((2, 1), 1.0)) + for species_idx in range(3): + assert ff._get_ctx_val(d, "mass", mass=7.0, + species_idx=species_idx) == 7.0 + + def test_per_species_array_is_picked_by_species_index(self): + d = _field(np.full((2, 1), 1.0)) + for species_idx, expected in enumerate([1.0, 2.0, 3.0]): + got = ff._get_ctx_val(d, + "mass", + mass=[1.0, 2.0, 3.0], + species_idx=species_idx) + assert got == expected + + def test_array_without_a_species_index_is_an_error(self): + d = _field(np.full((2, 1), 1.0)) + with pytest.raises(KeyError, match="not resolved per species"): + ff._get_ctx_val(d, "mass", mass=[1.0, 2.0]) + + def test_too_short_an_array_is_an_error(self): + d = _field(np.full((2, 1), 1.0)) + with pytest.raises(ValueError, match="only 2 values"): + ff._get_ctx_val(d, "mass", mass=[1.0, 2.0], species_idx=2, species="ion2") + + def test_missing_everywhere_is_an_error(self): + d = _field(np.full((2, 1), 1.0)) + with pytest.raises(KeyError, match="mass"): + ff._get_ctx_val(d, "mass") + + +class TestHeatFluxes: + """Lab-frame energy fluxes and fluid-frame heat fluxes.""" + + def test_qpar_lab_frame(self): + m3par = _field(np.full((2, 1), 4.0), mass=2.0) + out = ff.fetch_qpar([m3par]) + np.testing.assert_allclose(out.values[..., 0], 0.5 * 2.0 * 4.0) + + def test_qperp_lab_frame(self): + m3perp = _field(np.full((2, 1), 6.0), mass=3.0) + out = ff.fetch_qperp([m3perp]) + np.testing.assert_allclose(out.values[..., 0], 0.5 * 3.0 * 6.0) + + def test_qpar_fluid_matches_the_hand_derived_formula(self): + m0 = _field(np.full((2, 1), 2.0), mass=5.0) + m1 = _field(np.full((2, 1), 6.0)) + m2par = _field(np.full((2, 1), 10.0)) + m3par = _field(np.full((2, 1), 40.0)) + out = ff.fetch_qpar_fluid([m0, m1, m2par, m3par]) + upar = 6.0 / 2.0 + expected = 0.5 * 5.0 * (40.0 - 3.0 * upar * 10.0 + 2.0 * upar**2 * 6.0) + np.testing.assert_allclose(out.values[..., 0], expected) + + def test_qperp_fluid_matches_the_hand_derived_formula(self): + m0 = _field(np.full((2, 1), 2.0), mass=5.0) + m1 = _field(np.full((2, 1), 6.0)) + m2perp = _field(np.full((2, 1), 8.0)) + m3perp = _field(np.full((2, 1), 20.0)) + out = ff.fetch_qperp_fluid([m0, m1, m2perp, m3perp]) + upar = 6.0 / 2.0 + expected = 0.5 * 5.0 * (20.0 - upar * 8.0) + np.testing.assert_allclose(out.values[..., 0], expected) + + def test_qpar_fluid_vanishes_for_a_maxwellian(self): + """A Maxwellian carries no parallel heat flux in the fluid frame: the + three terms of ``(m/2)*[M3par - 3*u*M2par + 2*u^2*M1]`` cancel exactly + for ``M1=n*u``, ``M2par=n*(u^2+T/m)``, ``M3par=n*(u^3+3*u*T/m)``, so the + residual is compared against the size of an individual term rather + than an absolute zero.""" + n, u, T, m = 2.7e19, 1.3e4, 9.5e-18, 3.343e-27 + vt_sq = T / m + m0 = _field(np.full((2, 1), n), mass=m) + m1 = _field(np.full((2, 1), n * u)) + m2par = _field(np.full((2, 1), n * (u**2 + vt_sq))) + m3par = _field(np.full((2, 1), n * (u**3 + 3.0 * u * vt_sq))) + out = ff.fetch_qpar_fluid([m0, m1, m2par, m3par]) + term_scale = 0.5 * m * n * abs(u)**3 + assert np.all(np.abs(out.values[..., 0]) / term_scale < 1e-9) + + +class TestThermalSpeedAndLengths: + """``vt``/Larmor-radius/Debye-length -- plain ``np.sqrt`` on interpolated + data (this layer never touches ``dg``/``gpython``; see the module + docstring in ``quantities.py``).""" + + def test_vt(self): + temp = _field(np.full((2, 1), 8.0), mass=2.0) + out = ff.fetch_vt([temp]) + np.testing.assert_allclose(out.values[..., 0], np.sqrt(4.0)) + + def test_larmor_radius(self): + temp = _field(np.full((2, 1), 4.0), mass=9.0, charge=-2.0) + bmag = _field(np.full((2, 1), 3.0)) + out = ff.fetch_larmor_radius([temp, bmag]) + expected = np.sqrt(9.0 * 4.0) / (2.0 * 3.0) + np.testing.assert_allclose(out.values[..., 0], expected) + + def test_debye_length(self): + from scipy import constants + temp = _field(np.full((2, 1), 5.0), charge=2.0) + m0 = _field(np.full((2, 1), 7.0)) + out = ff.fetch_debye_length([temp, m0]) + expected = np.sqrt(constants.epsilon_0 * 5.0 / (7.0 * 4.0)) + np.testing.assert_allclose(out.values[..., 0], expected) + + +class TestSoundSpeed: + """The multi-species sound speeds, dispatched by '--extra kind='. + + ``fetch_c_s`` is an ``is_multi_species`` fetch function: it receives one + ``[M0, temp]`` source list per species (as + ``GkQuantity.fetch_multi``/``load_quantity`` would hand it), not a + flat list. + """ + + @staticmethod + def _species(dens, temp, mass, charge): + return [ + _field(np.full((2, 1), dens), mass=mass, charge=charge), + _field(np.full((2, 1), temp), mass=mass, charge=charge) + ] + + def test_ion_acoustic_single_ion_species(self): + """With one Z=1 ion species the formula collapses to sqrt(Te/mi).""" + from scipy import constants + e = constants.elementary_charge + m_e, T_e = constants.electron_mass, 9.5e-18 + n_i, T_i, m_i, z_i = 2.7e19, 6.1e-18, 3.343e-27, 1.0 + n_e = n_i * z_i + + out = ff.fetch_c_s([ + self._species(n_e, T_e, m_e, -e), + self._species(n_i, T_i, m_i, z_i * e) + ], + species=["elc", "ion"], + kind="ion_acoustic") + np.testing.assert_allclose(out.values[..., 0], + np.sqrt(T_e / m_i), + rtol=1e-10) + + def test_thermo_defaults_are_gamma_e_1_gamma_i_3(self): + from scipy import constants + e = constants.elementary_charge + m_e, T_e = constants.electron_mass, 9.5e-18 + n_i, T_i, m_i, z_i = 2.7e19, 6.1e-18, 3.343e-27, 1.0 + n_e = n_i * z_i + + gdatas = [ + self._species(n_e, T_e, m_e, -e), + self._species(n_i, T_i, m_i, z_i * e) + ] + default = ff.fetch_c_s(gdatas, species=["elc", "ion"], kind="thermo") + explicit = ff.fetch_c_s(gdatas, + species=["elc", "ion"], + kind="thermo", + gamma_e=1.0, + gamma_i=3.0) + np.testing.assert_allclose(default.values, explicit.values, rtol=1e-12) + + expected = np.sqrt((1.0 * n_e * T_e + 3.0 * n_i * T_i) / (n_i * m_i)) + np.testing.assert_allclose(default.values[..., 0], expected, rtol=1e-10) + + def test_species_order_does_not_matter(self): + """Species are identified by charge sign, so the order is irrelevant.""" + elc = self._species(1.0e19, 5.0e-18, 9.1e-31, -1.6e-19) + ion = self._species(1.0e19, 3.0e-18, 3.3e-27, 1.6e-19) + forward = ff.fetch_c_s([elc, ion], + species=["elc", "ion"], + kind="ion_acoustic") + shuffled = ff.fetch_c_s([ion, elc], + species=["ion", "elc"], + kind="ion_acoustic") + np.testing.assert_allclose(forward.values, shuffled.values, rtol=1e-12) + + def test_no_electron_species_is_an_error(self): + ion = self._species(1.0, 1.0, 1.0, 1.0) + with pytest.raises(ValueError, match="exactly one negatively charged"): + ff.fetch_c_s([ion], species=["ion"]) + + def test_no_ion_species_is_an_error(self): + elc = self._species(1.0, 1.0, 1.0, -1.0) + with pytest.raises(ValueError, match="no positively charged"): + ff.fetch_c_s([elc], species=["elc"]) + + def test_unknown_kind_is_an_error(self): + elc = self._species(1.0, 1.0, 1.0, -1.0) + ion = self._species(1.0, 1.0, 1.0, 1.0) + with pytest.raises(ValueError, match="unknown kind"): + ff.fetch_c_s([elc, ion], species=["elc", "ion"], kind="bogus") + + def test_per_species_extra_arrays_reach_nested_sources(self): + """'--extra mass=1,2,charge=-1,1' must give each species its own entry, + threaded through even though these sources carry no mass/charge in + their own ctx -- the whole point of ``species_idx`` reaching + ``get_src_gdata``/``_split_elc_ions``.""" + bare = lambda dens, temp: [ + _field(np.full((2, 1), dens)), + _field(np.full((2, 1), temp)) + ] + out = ff.fetch_c_s([bare(2.0, 8.0), bare(2.0, 8.0)], + species=["elc", "ion"], + kind="thermo", + mass=[1.0, 2.0], + charge=[-1.0, 1.0]) + # n=2, T=8 for both species; gamma_e=1, gamma_i=3 (defaults). + expected = np.sqrt((1.0 * 2.0 * 8.0 + 3.0 * 2.0 * 8.0) / (2.0 * 2.0)) + np.testing.assert_allclose(out.values[..., 0], expected, rtol=1e-10) + + +class TestNormalizedQuantities: + + def test_rho_over_lambda(self): + rho = _field(np.full((2, 1), 6.0)) + lambda_d = _field(np.full((2, 1), 3.0)) + out = ff.fetch_rho_over_lambda([rho, lambda_d]) + np.testing.assert_allclose(out.values[..., 0], 2.0) + + def test_phi_norm(self): + from scipy import constants + phi = _field(np.full((2, 1), 5.0)) + temp = _field(np.full((2, 1), 2.0)) + out = ff.fetch_phi_norm([phi, temp]) + np.testing.assert_allclose(out.values[..., 0], + constants.elementary_charge * 5.0 / 2.0) + + def test_qpar_norm(self): + q = _field(np.full((2, 1), 12.0)) + m0 = _field(np.full((2, 1), 2.0)) + temp = _field(np.full((2, 1), 3.0)) + c_s = _field(np.full((2, 1), 2.0)) + out = ff.fetch_qpar_norm([q, m0, temp, c_s]) + np.testing.assert_allclose(out.values[..., 0], 12.0 / (2.0 * 3.0 * 2.0)) + + +class TestDriftVelocities: + """``fetch_gradB_vel``/``fetch_diamag_vel`` and the remaining + ``_b_cross_grad_div_b_component`` branches (comp 1/2, cdim 1/2/3).""" + + def _synthetic(self, cdim, comp): + grid = [np.linspace(0.0, float(n), n + 1) for n in [4, 4, 4][:cdim]] + centers = [0.5 * (g[:-1] + g[1:]) for g in grid] + mesh = np.meshgrid(*centers, indexing="ij") + scalar = _field(sum(mesh)[..., np.newaxis], grid=grid) + jacobtot_inv = _field(np.full(scalar.values.shape, 2.0), grid=grid) + b_i = _field(np.stack([np.full(mesh[0].shape, float(k)) for k in range(3)], + axis=-1), + grid=grid) + return scalar, jacobtot_inv, b_i + + @pytest.mark.parametrize("cdim,comp", [(1, 0), (1, 1), (1, 2), (2, 0), (2, 1), + (2, 2), (3, 0), (3, 1), (3, 2)]) + def test_all_cdim_comp_combinations_run(self, cdim, comp): + scalar, jacobtot_inv, b_i = self._synthetic(cdim, comp) + out = ff._b_cross_grad_div_b_component(scalar, jacobtot_inv, b_i, comp) + assert out.values.shape == scalar.values.shape + assert np.all(np.isfinite(out.values)) + + def test_gradB_vel(self): + scalar, jacobtot_inv, b_i = self._synthetic(1, 0) + Tperp = _field(np.full(scalar.values.shape, 3.0), + grid=scalar.grid, + charge=2.0) + out = ff.fetch_gradB_vel([jacobtot_inv, scalar, b_i, Tperp], dir=0) + assert np.all(np.isfinite(out.values)) + + def test_diamag_vel(self): + scalar, jacobtot_inv, b_i = self._synthetic(1, 0) + m0 = _field(np.full(scalar.values.shape, 5.0), grid=scalar.grid) + pressperp = _field(np.full(scalar.values.shape, 3.0), + grid=scalar.grid, + charge=2.0) + out = ff.fetch_diamag_vel([jacobtot_inv, scalar, b_i, m0, pressperp], dir=0) + assert np.all(np.isfinite(out.values)) + + def test_gradB_vel_requires_dir(self): + with pytest.raises(KeyError): + ff.fetch_gradB_vel([None, None, None, None]) + + def test_diamag_vel_requires_dir(self): + with pytest.raises(KeyError): + ff.fetch_diamag_vel([None, None, None, None, None]) + + +class TestLoadDistf: + """``fetch_funcs.load_distf`` -- the registry 'distf' quantity's fetch + function -- stubbed against ``load_distf`` so this checks the option + translation (``dict_get_bool``, path/name joining) without needing a real + distribution-function file set (covered end to end by + ``TestLoadGkDistfReal`` instead).""" + + def test_forwards_options(self, monkeypatch): + calls = {} + + def fake_load_distf(**kwargs): + calls.update(kwargs) + return "sentinel" + + from postgkyl.diagnostics.gk import distf as distf_mod + monkeypatch.setattr(distf_mod, "load_distf", fake_load_distf) + + out = ff.load_distf([], + path="/some/path/", + name="sim", + species="ion", + frame="3", + suffix="src", + c2p_vel="0", + mc2nu="1", + block=2) + assert out == "sentinel" + assert calls["name"] == "/some/path/sim" + assert calls["species"] == "ion" + assert calls["frame"] == 3 + assert calls["suffix"] == "src" + assert calls["use_c2p_vel"] is False + assert calls["use_mc2nu"] is True + assert calls["use_mapc2p"] is False + assert calls["block_idx"] == 2 + assert calls["num_interp"] == 0 + + def test_defaults(self, monkeypatch): + calls = {} + + def fake_load_distf(**kwargs): + calls.update(kwargs) + return "sentinel" + + from postgkyl.diagnostics.gk import distf as distf_mod + monkeypatch.setattr(distf_mod, "load_distf", fake_load_distf) + + ff.load_distf([], path="p", name="n", species="ion", frame=0) + # c2p_vel defaults True when not given as an extra. + assert calls["use_c2p_vel"] is True + + +class TestCrossGradDivB: + """``_b_cross_grad_div_b_component`` on a 1-D synthetic field (cdim=1): + only the 'positive' term is defined, so the formula reduces to + ``d(f)/dx * b_i[bi_c_pos] * jacobtot_inv``.""" + + def test_linear_scalar_1d(self): + x = np.linspace(0.0, 4.0, 5) # 4 cells, dx=1 + centers = 0.5 * (x[:-1] + x[1:]) # phi(x) = x at cell centers + phi = _field(centers[:, np.newaxis], grid=[x]) + jacobtot_inv = _field(np.full((4, 1), 2.0), grid=[x]) + b_i = _field(np.tile([0.0, 1.0, 0.0], (4, 1)), grid=[x]) + out = ff._b_cross_grad_div_b_component(phi, jacobtot_inv, b_i, 0) + np.testing.assert_allclose(out.values[..., 0], 2.0, rtol=1e-6) + + def test_invalid_component_raises(self): + x = np.linspace(0.0, 1.0, 3) + phi = _field(np.zeros((2, 1)), grid=[x]) + jacobtot_inv = _field(np.ones((2, 1)), grid=[x]) + b_i = _field(np.zeros((2, 3)), grid=[x]) + with pytest.raises(KeyError): + ff._b_cross_grad_div_b_component(phi, jacobtot_inv, b_i, 3) + + def test_ExB_vel_requires_dir(self): + with pytest.raises(KeyError): + ff.fetch_ExB_vel([None, None, None, None]) + + +class TestLoadQuantity: + + def test_available_quantities_sorted(self): + names = available_quantities() + assert names == sorted(names) + assert "M0" in names + assert "distf" in names + + def test_unknown_quantity_raises(self): + with pytest.raises(ValueError, match="Unknown quantity"): + load_quantity("not_a_quantity", None, "sim", path=DATA) + + @needs_gkeyll + def test_M0_from_hamiltonian_moments_real(self): + out = load_quantity("M0", "ion", HMOM_NAME, "250", path=DATA) + assert len(out) == 1 + assert out[0].get_label() == r"$M_{0i}$ (m$^{-3}$)" + assert out[0].values.shape[-1] == 1 + + @needs_gkeyll + def test_M1_from_hamiltonian_moments_real(self): + out = load_quantity("M1", "ion", HMOM_NAME, "250", path=DATA, mass=2.0) + assert len(out) == 1 + assert np.all(np.isfinite(out[0].values)) + + @needs_gkeyll + def test_geo_quantity_real(self): + out = load_quantity("geo_int_jacobtot_inv", None, GK_NAME, path=DATA) + assert len(out) == 1 + assert out[0].get_label() == r"$(J B)^{-1}$" + + @needs_gkeyll + def test_geo_quantity_missing_file_raises(self): + with pytest.raises(FileNotFoundError): + load_quantity("geo_int_bmag", None, GK_NAME, path=DATA) + + def test_label_and_tag_override(self, tmp_path, monkeypatch): + # A species-independent geo quantity needs only its own marker file. + (tmp_path / "sim-geo_int_bmag.gkyl").touch() + monkeypatch.setattr(qmod, "GData", + lambda *a, **k: _field(np.full((2, 1), 3.0))) + out = load_quantity("geo_int_bmag", + None, + "sim", + path=str(tmp_path), + tag="mytag", + label="custom") + assert out[0].get_tag() == "mytag" + assert out[0].get_label() == "custom" + + +class _SyntheticSource: + """Serves a small, self-consistent constant-valued synthetic DG dataset + for every source file a quantity asks for -- ported from + tests_bak/test_gk_load_quantity.py's ``_make_synthetic_gdata``, adapted to + push through the new ``GDataState``/``.interpolate()`` (no ``ctypes``). + + Every source is served the same synthetic values; only the charge is read + back out of the file name (negative for an ``elc`` species, per + ``_ELC_SPECIES``), so multi-species quantities -- which tell electrons + from ions by the sign of the charge -- see a genuine electron species. + """ + + POLY_ORDER = 1 + BASIS_TYPE = "serendipity" + NUM_BASIS = 2 + NUM_PHYS_COMPS = 4 + NUM_CELLS = 4 + + def __call__(self, *args, **kwargs): + values = np.zeros((self.NUM_CELLS, self.NUM_BASIS * self.NUM_PHYS_COMPS)) + for comp in range(self.NUM_PHYS_COMPS): + values[:, comp * self.NUM_BASIS] = (comp + 2) * np.sqrt(2.0) + file_name = str(args[0]) if args else "" + charge = -1.0 if f"-{_ELC_SPECIES}_" in file_name else 1.0 + grid = [np.linspace(0.0, 1.0, self.NUM_CELLS + 1)] + d = GDataState( + ctx={ + "poly_order": self.POLY_ORDER, + "basis_type": self.BASIS_TYPE, + "mass": 1.0, + "charge": charge + }) + d.push(grid, values) + return d + + +# Species names used to drive a genuine electron/ion split in the synthetic +# smoke tests: multi-species quantities (e.g. the sound speed) tell them +# apart by the sign of the charge, which _SyntheticSource keys off this name. +_ELC_SPECIES = "elc" +_ION_SPECIES = "ion" + + +def _collect_source_files(quant, path, name, species, frame) -> set: + files: set[str] = set() + for combo in quant.source: + for src in combo: + if isinstance(src, str): + files.add(quant._src_file_name(path, name, species, src, frame)) + else: + files |= _collect_source_files(src, path, name, species, frame) + return files + + +def _extra_for(quant) -> dict: + extra = {} + if quant.is_vector: + extra["direction"] = 0 + return extra + + +@needs_gkeyll +@pytest.mark.parametrize("quantity", gk_quant_registry.list()) +def test_every_registered_quantity_produces_a_dataset(quantity, tmp_path, + monkeypatch): + """Smoke test across the whole registry (weak assertion, matching + tests_bak/test_gk_load_quantity.py): the synthetic data is not physically + consistent across different marker files (every file gets the SAME + constant recipe, regardless of what real quantity it names), so this + checks "no exception, one dataset comes back", not specific numbers -- + those are covered analytically in ``TestFetchPhysics`` above.""" + if quantity == "distf": + pytest.skip("distf delegates to load_distf, covered by " + "TestLoadGkDistfReal against the real staged fixtures") + + quant = gk_quant_registry.get(quantity) + name, frame = "gktest", 0 + path = str(tmp_path) + # A multi-species quantity (e.g. the sound speed) needs an electron and + # an ion species to combine; every other quantity is fine with just one. + species = (f"{_ELC_SPECIES},{_ION_SPECIES}" + if quant.is_multi_species else _ION_SPECIES) + + for species_name in species.split(","): + for file_name in _collect_source_files(quant, path, name, species_name, + frame): + open(file_name, "w").close() + + monkeypatch.setattr(qmod, "GData", _SyntheticSource()) + + out = load_quantity(quantity, + species, + name, + str(frame), + path=path, + **_extra_for(quant)) + assert len(out) >= 1 + assert isinstance(out[0], GDataState) + + +class TestGkQuantityGetAvailSource: + """``GkQuantity.get_avail_source``/``_avail_combo_frames`` frame-list + parsing branches, exercised directly (rather than through the full + registry) for precise control over which frames each source combo has.""" + + def _touch_frames(self, tmp_path, stem, frames): + for f in frames: + (tmp_path / f"{stem}{f}.gkyl").touch() + + def test_comma_separated_frame_list(self, tmp_path): + quant = qmod.GkQuantity(name="q", + source=[["a"]], + fetch_func=[None], + label="q", + is_species_dep=True) + self._touch_frames(tmp_path, "sim-ion_a_", [0, 2, 4]) + combo_idx, frames = quant.get_avail_source(str(tmp_path), "sim", "ion", + "0,2") + assert combo_idx == 0 + assert frames == [0, 2] + + def test_none_frame_selects_every_available(self, tmp_path): + quant = qmod.GkQuantity(name="q", + source=[["a"]], + fetch_func=[None], + label="q", + is_species_dep=True) + self._touch_frames(tmp_path, "sim-ion_a_", [0, 1, 3]) + combo_idx, frames = quant.get_avail_source(str(tmp_path), "sim", "ion", + None) + assert frames == [0, 1, 3] + + def test_partial_range_frame(self, tmp_path): + quant = qmod.GkQuantity(name="q", + source=[["a"]], + fetch_func=[None], + label="q", + is_species_dep=True) + self._touch_frames(tmp_path, "sim-ion_a_", [0, 1, 2, 3]) + combo_idx, frames = quant.get_avail_source(str(tmp_path), "sim", "ion", + "1:") + assert frames == [1, 2, 3] + + def test_mismatched_frame_sets_falls_back_to_next_combo(self, tmp_path): + # combo 0 ("a","b") has mismatched frame sets -> rejected; combo 1 ("c") + # is used instead. + quant = qmod.GkQuantity(name="q", + source=[["a", "b"], ["c"]], + fetch_func=[None, None], + label="q", + is_species_dep=True) + self._touch_frames(tmp_path, "sim-ion_a_", [0, 1]) + self._touch_frames(tmp_path, "sim-ion_b_", [0]) + self._touch_frames(tmp_path, "sim-ion_c_", [5]) + combo_idx, frames = quant.get_avail_source(str(tmp_path), "sim", "ion", + None) + assert combo_idx == 1 + assert frames == [5] + + def test_no_files_found_raises(self, tmp_path): + quant = qmod.GkQuantity(name="q", + source=[["a"]], + fetch_func=[None], + label="q", + is_species_dep=True) + with pytest.raises(FileNotFoundError): + quant.get_avail_source(str(tmp_path), "sim", "ion", None) + + +@needs_gkeyll +class TestLoadQuantityMultiSpeciesMultiFrame: + """Exercises ``load_quantity``'s multi-species/multi-frame label/tag + suffix branches (only reached when more than one species or frame is + requested).""" + + def test_multiple_species(self, tmp_path, monkeypatch): + quant = gk_quant_registry.get("M0") + name = "gktest" + path = str(tmp_path) + for species in ("ion", "elc"): + for file_name in _collect_source_files(quant, path, name, species, 0): + open(file_name, "w").close() + monkeypatch.setattr(qmod, "GData", _SyntheticSource()) + + out = load_quantity("M0", + "ion,elc", + name, + "0", + path=path, + tag="t", + label="custom") + assert len(out) == 2 + assert {d.get_tag() for d in out} == {"t_ion", "t_elc"} + assert {d.get_label() for d in out} == {"custom ion", "custom elc"} + + def test_multiple_frames_suffixes_label(self, tmp_path, monkeypatch): + quant = gk_quant_registry.get("M0") + name = "gktest" + path = str(tmp_path) + for frame in (0, 1, 2): + for file_name in _collect_source_files(quant, path, name, "ion", frame): + open(file_name, "w").close() + monkeypatch.setattr(qmod, "GData", _SyntheticSource()) + + out = load_quantity("M0", "ion", name, None, path=path) + assert len(out) == 3 + assert all(" f" in d.get_label() for d in out) + + def test_multi_species_quantity_yields_a_single_dataset( + self, tmp_path, monkeypatch): + """A multi-species quantity (the sound speed) combines its species into + one dataset, unlike a per-species quantity (M0), which still produces + one dataset per species.""" + name = "gktest" + path = str(tmp_path) + for species in (_ELC_SPECIES, _ION_SPECIES): + for file_name in _collect_source_files(gk_quant_registry.get("c_s"), path, + name, species, 0): + open(file_name, "w").close() + for file_name in _collect_source_files(gk_quant_registry.get("M0"), path, + name, species, 0): + open(file_name, "w").close() + monkeypatch.setattr(qmod, "GData", _SyntheticSource()) + + species = f"{_ELC_SPECIES},{_ION_SPECIES}" + out = load_quantity("c_s", species, name, "0", path=path) + assert len(out) == 1 + + out = load_quantity("M0", species, name, "0", path=path) + assert len(out) == 2 + + def test_multi_species_quantity_needs_a_species_list(self, tmp_path): + with pytest.raises(ValueError, match="needs a species list"): + load_quantity("c_s", None, "gktest", "0", path=str(tmp_path)) + + +class TestUtils: + """postgkyl.diagnostics.gk.utils -- file/geometry helpers ported + from src_bak's gk_utils.py (matplotlib bits dropped, read_g*file adapted + to postgkyl.gdata.load + .interpolate()).""" + + def test_dict_get_bool_default(self): + assert utils.dict_get_bool({}, "k", True) is True + assert utils.dict_get_bool({}, "k", False) is False + + def test_dict_get_bool_string_true_variants(self): + assert utils.dict_get_bool({"k": "1"}, "k", False) is True + assert utils.dict_get_bool({"k": "True"}, "k", False) is True + assert utils.dict_get_bool({"k": " true "}, "k", False) is True + + def test_dict_get_bool_string_false(self): + assert utils.dict_get_bool({"k": "0"}, "k", True) is False + assert utils.dict_get_bool({"k": "no"}, "k", True) is False + + def test_dict_get_bool_non_string(self): + assert utils.dict_get_bool({"k": 1}, "k", False) is True + assert utils.dict_get_bool({"k": 0}, "k", True) is False + + def test_parse_slice_string(self): + assert utils.parse_slice_string("1:5") == slice(1, 5) + assert utils.parse_slice_string(":5") == slice(None, 5) + assert utils.parse_slice_string("1:") == slice(1, None) + assert utils.parse_slice_string("1:5:2") == slice(1, 5, 2) + + def test_parse_slice_string_invalid_raises(self): + with pytest.raises(ValueError): + utils.parse_slice_string("a:5") + + def test_get_block_indices_single(self): + assert utils.get_block_indices("-10", "unused") == [0] + + def test_get_block_indices_all(self, tmp_path): + for i in range(3): + (tmp_path / f"sim_b{i}-ion_field_0.gkyl").touch() + pattern = str(tmp_path / "sim_b*-ion_field_0.gkyl") + assert utils.get_block_indices("-1", pattern) == [0, 1, 2] + + def test_get_block_indices_comma_list(self): + assert utils.get_block_indices("0,2,4", "unused") == [0, 2, 4] + + def test_get_block_indices_slice(self): + assert utils.get_block_indices("1:4", "unused") == [1, 2, 3] + + def test_get_block_indices_single_int(self): + assert utils.get_block_indices("2", "unused") == [2] + + def test_get_block_indices_invalid_raises(self): + with pytest.raises(NameError): + utils.get_block_indices("not-a-spec", "unused") + + @needs_gkeyll + def test_read_gfile(self): + grid, values, gdata = utils.read_gfile( + os.path.join(DATA, f"{GK_NAME}-geo_int_jacobtot_inv.gkyl")) + assert values.shape[0] == gdata.num_cells[0] + + def test_read_gfile_if_present_missing(self, tmp_path): + found, grid, values, gdata = utils.read_gfile_if_present( + str(tmp_path / "does_not_exist.gkyl")) + assert found is False + assert grid is None and values is None and gdata is None + + @needs_gkeyll + def test_read_gfile_if_present_found(self): + found, grid, values, gdata = utils.read_gfile_if_present( + os.path.join(DATA, f"{GK_NAME}-geo_int_jacobtot_inv.gkyl")) + assert found is True + assert values is not None + + @needs_gkeyll + def test_read_interpolated_gfile(self): + grid, values, gdata = utils.read_interpolated_gfile( + os.path.join(DATA, f"{GK_NAME}-geo_int_jacobtot_inv.gkyl"), + poly_order=1, + basis_type="serendipity") + assert gdata.is_interpolated + + @needs_gkeyll + def test_read_interpolated_gfile_with_comp(self): + grid, values, gdata = utils.read_interpolated_gfile( + os.path.join(DATA, f"{GK_NAME}-geo_int_jacobtot_inv.gkyl"), + poly_order=1, + basis_type="serendipity", + comp=0) + assert gdata.num_comps == 1 diff --git a/tests/test_diagnostics_kinetic.py b/tests/test_diagnostics_kinetic.py new file mode 100644 index 00000000..2fdd49ef --- /dev/null +++ b/tests/test_diagnostics_kinetic.py @@ -0,0 +1,159 @@ +"""Tests for postgkyl.diagnostics.vm.kinetic -- distribution-function frame +transform, folding the array-math analytic tests (formerly +tests_models_frame.py) with the verb-level guard/inplace tests (formerly +part of tests_ops_physics.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.vm import kinetic +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +class TestTransformFrameCdim1: + + def _distribution(self, nx=2, nv=3): + x_edges = np.linspace(0.0, 1.0, nx + 1) + v_edges = np.linspace(-2.0, 2.0, nv + 1) + values = np.ones((nx, nv, 1)) + return _make([x_edges, v_edges], values) + + def test_basic_returns_unchanged_values(self): + f = _make([np.linspace(0.0, 1.0, 4), + np.linspace(-3.0, 3.0, 5)], np.ones((3, 4, 1))) + bulk = _make([f.grid[0]], np.ones((3, 1)) * 0.5) + out = kinetic.transform_frame(f, bulk, cdim=1) + np.testing.assert_array_equal(out.values, f.values) + assert len(out.grid) == 2 + + def test_zero_velocity_leaves_grid_unshifted(self): + v_grid = np.linspace(-2.0, 2.0, 4) + f = _make([np.linspace(0.0, 1.0, 3), v_grid], + np.random.default_rng(0).random((2, 3, 1))) + bulk = _make([f.grid[0]], np.zeros((2, 1))) + out = kinetic.transform_frame(f, bulk, cdim=1) + np.testing.assert_array_equal(out.values, f.values) + np.testing.assert_allclose(out.grid[1], np.tile(v_grid, (3, 1))) + + def test_shifts_velocity_grid_by_bulk_velocity(self): + v_grid = np.linspace(-2.0, 2.0, 4) + f = _make([np.linspace(0.0, 1.0, 3), v_grid], np.ones((2, 3, 1))) + bulk = _make([f.grid[0]], np.full((2, 1), 0.5)) + out = kinetic.transform_frame(f, bulk, cdim=1) + # Interior nodes see the average of the two neighboring cells' shift + # (both 0.5 here); edge nodes see the single adjacent cell's shift. + np.testing.assert_allclose(out.grid[1][0], v_grid + 0.5) + np.testing.assert_allclose(out.grid[1][-1], v_grid + 0.5) + + def test_matches_private_helper(self): + f = self._distribution() + bulk = _make([f.grid[0]], np.array([[0.1], [0.2]])) + out = kinetic.transform_frame(f, bulk, cdim=1) + grid, values = kinetic._transform_frame(f.grid, f.values, bulk.values, 1) + for d in range(2): + np.testing.assert_allclose(out.grid[d], grid[d]) + np.testing.assert_allclose(out.values, values) + + def test_inplace_mutates_distribution(self): + f = self._distribution() + bulk = _make([f.grid[0]], np.array([[0.1], [0.2]])) + out = kinetic.transform_frame(f, bulk, cdim=1, inplace=True) + assert out is f + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + bulk = _make([np.array([0.0, 1.0])], np.array([[0.1]])) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + kinetic.transform_frame(d, bulk, cdim=1) + + +class TestTransformFrameCdim2: + + def test_zero_velocity_leaves_grid_unshifted(self): + nx, ny, nv = 2, 2, 3 + x_grid = np.linspace(0.0, 1.0, nx + 1) + y_grid = np.linspace(0.0, 1.0, ny + 1) + grid_f = [x_grid, y_grid, np.linspace(-2.0, 2.0, nv + 1)] + values_f = np.ones((nx, ny, nv, 1)) + f = _make(grid_f, values_f) + bulk = _make([f.grid[0], f.grid[1]], np.zeros((nx, ny, 1))) + out = kinetic.transform_frame(f, bulk, cdim=2) + np.testing.assert_array_equal(out.values, values_f) + assert len(out.grid) == 3 + np.testing.assert_allclose(out.grid[2], + np.tile(grid_f[2], (nx + 1, ny + 1, 1))) + + def test_shifts_velocity_grid_by_bulk_velocity(self): + nx, ny, nv = 2, 2, 3 + v_grid = np.linspace(-2.0, 2.0, nv + 1) + grid_f = [ + np.linspace(0.0, 1.0, nx + 1), + np.linspace(0.0, 1.0, ny + 1), v_grid + ] + values_f = np.ones((nx, ny, nv, 1)) + f = _make(grid_f, values_f) + bulk = _make([f.grid[0], f.grid[1]], np.full((nx, ny, 1), 0.5)) + out = kinetic.transform_frame(f, bulk, cdim=2) + np.testing.assert_array_equal(out.values, values_f) + # Every corner node sees the same 0.5 shift, since the bulk velocity is + # uniform. + np.testing.assert_allclose(out.grid[2][0, 0], v_grid + 0.5) + np.testing.assert_allclose(out.grid[2][-1, -1], v_grid + 0.5) + + +class TestTransformFrameCdim3: + + def test_zero_velocity_leaves_grid_unshifted(self): + nx, ny, nz, nv = 2, 2, 2, 2 + grid_f = [ + np.linspace(0.0, 1.0, nx + 1), + np.linspace(0.0, 1.0, ny + 1), + np.linspace(0.0, 1.0, nz + 1), + np.linspace(-2.0, 2.0, nv + 1) + ] + values_f = np.ones((nx, ny, nz, nv, 1)) + f = _make(grid_f, values_f) + bulk = _make([f.grid[0], f.grid[1], f.grid[2]], np.zeros((nx, ny, nz, 1))) + out = kinetic.transform_frame(f, bulk, cdim=3) + np.testing.assert_array_equal(out.values, values_f) + assert len(out.grid) == 4 + np.testing.assert_allclose(out.grid[3], + np.tile(grid_f[3], (nx + 1, ny + 1, nz + 1, 1))) + + def test_shifts_velocity_grid_by_bulk_velocity(self): + nx, ny, nz, nv = 2, 2, 2, 2 + v_grid = np.linspace(-2.0, 2.0, nv + 1) + grid_f = [ + np.linspace(0.0, 1.0, nx + 1), + np.linspace(0.0, 1.0, ny + 1), + np.linspace(0.0, 1.0, nz + 1), v_grid + ] + values_f = np.ones((nx, ny, nz, nv, 1)) + f = _make(grid_f, values_f) + bulk = _make([f.grid[0], f.grid[1], f.grid[2]], np.full((nx, ny, nz, 1), + 0.5)) + out = kinetic.transform_frame(f, bulk, cdim=3) + np.testing.assert_array_equal(out.values, values_f) + np.testing.assert_allclose(out.grid[3][0, 0, 0], v_grid + 0.5) + np.testing.assert_allclose(out.grid[3][-1, -1, -1], v_grid + 0.5) diff --git a/tests/test_diagnostics_mhd.py b/tests/test_diagnostics_mhd.py new file mode 100644 index 00000000..1c099d6f --- /dev/null +++ b/tests/test_diagnostics_mhd.py @@ -0,0 +1,143 @@ +"""Tests for postgkyl.diagnostics.mom.mhd -- MHD B-field, pressure, temperature, +sound speed, Mach number, folding the array-math analytic tests (formerly +tests_models_mhd.py) with the verb-level guard/VARIABLES tests (formerly +part of tests_ops_moments.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.mom import mhd +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +_G1D = [np.array([0.0, 1.0])] + +_RHO = 1.0 +_VX = 0.5 +_P_THERMAL = 0.6 +_GAMMA = 5.0 / 3.0 +_BX, _BY, _BZ = 1.0, 0.0, 0.0 +_MAG_P = 0.5 * (_BX**2 + _BY**2 + _BZ**2) +_E_MHD = 0.5 * _RHO * _VX**2 + _P_THERMAL / (_GAMMA - 1) + _MAG_P +_MHD8 = np.array([[_RHO, _RHO * _VX, 0.0, 0.0, _E_MHD, _BX, _BY, _BZ]]) + + +class TestFieldExtraction: + + def test_bx(self): + d = _make(_G1D, _MHD8) + np.testing.assert_allclose(mhd.bx(d).values[0, 0], _BX) + + def test_by(self): + d = _make(_G1D, _MHD8) + np.testing.assert_allclose(mhd.by(d).values[0, 0], _BY) + + def test_bz(self): + d = _make(_G1D, _MHD8) + np.testing.assert_allclose(mhd.bz(d).values[0, 0], _BZ) + + def test_bi_shape_and_values(self): + d = _make(_G1D, _MHD8) + out = mhd.bi(d) + assert out.values.shape[-1] == 3 + np.testing.assert_allclose(out.values[0], [_BX, _BY, _BZ]) + + def test_mag_pressure(self): + d = _make(_G1D, _MHD8) + out = mhd.mag_pressure(d) + np.testing.assert_allclose(out.values[0, 0], _MAG_P) + + @needs_gkeyll + def test_bx_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + mhd.bx(d) + + +class TestThermo: + + def test_pressure(self): + d = _make(_G1D, _MHD8) + out = mhd.pressure(d) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL, rtol=1e-10) + + def test_temp(self): + d = _make(_G1D, _MHD8) + out = mhd.temp(d) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL / _RHO, rtol=1e-10) + + def test_sound(self): + d = _make(_G1D, _MHD8) + out = mhd.sound(d) + expected = np.sqrt(_GAMMA * _P_THERMAL / _RHO) + np.testing.assert_allclose(out.values[0, 0], expected, rtol=1e-10) + + def test_mach(self): + d = _make(_G1D, _MHD8) + out = mhd.mach(d) + cs = np.sqrt(_GAMMA * _P_THERMAL / _RHO) + np.testing.assert_allclose(out.values[0, 0], _VX / cs, rtol=1e-10) + + def test_mag_p_zero_field_gives_pure_gas_pressure(self): + e = _P_THERMAL / (_GAMMA - 1) + 0.5 * _RHO * _VX**2 + values = np.array([[_RHO, _RHO * _VX, 0.0, 0.0, e, 0.0, 0.0, 0.0]]) + d = _make(_G1D, values) + out = mhd.pressure(d) + np.testing.assert_allclose(out.values[0, 0], _P_THERMAL, rtol=1e-10) + + def test_mu_0_is_forwarded_to_mag_pressure(self): + d = _make(_G1D, _MHD8) + out = mhd.mag_pressure(d, mu_0=2.0) + np.testing.assert_allclose(out.values[0, 0], _MAG_P / 2.0) + + @needs_gkeyll + def test_pressure_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + mhd.pressure(d) + + +class TestFiveMomentSetReused: + + def test_density_xvel_reused_from_five_moment(self): + from postgkyl.diagnostics.mom import five_moment as fm + assert mhd.density is fm.density + assert mhd.xvel is fm.xvel + assert mhd.yvel is fm.yvel + assert mhd.zvel is fm.zvel + assert mhd.vel is fm.vel + + +class TestVariables: + + def test_variables_table_has_exactly_the_old_mhd_vocabulary(self): + assert set(mhd.VARIABLES) == { + "density", "xvel", "yvel", "zvel", "vel", "Bx", "By", "Bz", "Bi", + "magpressure", "pressure", "temp", "sound", "mach" + } + + def test_variables_table_maps_to_public_functions(self): + assert mhd.VARIABLES["Bx"] is mhd.bx + assert mhd.VARIABLES["Bi"] is mhd.bi + assert mhd.VARIABLES["magpressure"] is mhd.mag_pressure + assert mhd.VARIABLES["density"] is mhd.density diff --git a/tests/test_diagnostics_multispecies.py b/tests/test_diagnostics_multispecies.py new file mode 100644 index 00000000..e2a5d248 --- /dev/null +++ b/tests/test_diagnostics_multispecies.py @@ -0,0 +1,170 @@ +"""Tests for postgkyl.diagnostics.mom.multispecies -- energy decomposition and +current accumulation, folding the array-math analytic tests (formerly +tests_models_energetics.py) with the verb-level guard/inplace tests +(formerly part of tests_ops_physics.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.mom import multispecies as ms +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + +_G1D = [np.array([0.0, 1.0])] +_GAMMA = 5.0 / 3.0 + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +def _make_5mom(rho, vx, p): + E = p / (_GAMMA - 1) + 0.5 * rho * vx**2 + return _make(_G1D, np.array([[rho, rho * vx, 0.0, 0.0, E]])) + + +class TestEnergetics: + + def test_components_and_total(self): + elc = _make_5mom(rho=1.0, vx=1.0, p=0.3) + ion = _make_5mom(rho=1.0, vx=0.5, p=0.6) + field = _make(_G1D, np.array([[1.0, 0.0, 0.0, 2.0, 0.0, + 0.0]])) # Ex=1, Bx=2 + + out = ms.energetics(elc, ion, field) + + assert out.values.shape[-1] == 7 + pre_expected = 0.3 + kee_expected = 0.5 * 1.0 * 1.0**2 + pri_expected = 0.6 + kei_expected = 0.5 * 1.0 * 0.5**2 + esq_expected = 1.0**2 / 2.0 + bsq_expected = 2.0**2 / 2.0 + np.testing.assert_allclose(out.values[0, 0], pre_expected, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 1], kee_expected, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 2], pri_expected, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 3], kei_expected, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 4], esq_expected, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 5], bsq_expected, rtol=1e-10) + total = (pre_expected + kee_expected + pri_expected + kei_expected + + esq_expected + bsq_expected) + np.testing.assert_allclose(out.values[0, 6], total, rtol=1e-10) + + def test_result_carries_field_grid(self): + elc = _make_5mom(rho=1.0, vx=0.0, p=1.0) + ion = _make_5mom(rho=1.0, vx=0.0, p=1.0) + field = _make(_G1D, np.zeros((1, 6))) + out = ms.energetics(elc, ion, field, inplace=True) + assert out is field + + def test_component_layout(self): + elc = _make_5mom(rho=1.0, vx=2.0, p=16.0 / 3.0) + ion = _make_5mom(rho=1.0, vx=2.0, p=16.0 / 3.0) + field = _make(_G1D, np.array([[1.0, 0.0, 0.0, 0.0, 2.0, 0.0]])) + out = ms.energetics(elc, ion, field) + comps = out.values[0] + np.testing.assert_allclose(comps[0], 16.0 / 3.0) # electron thermal + np.testing.assert_allclose(comps[1], 2.0) # electron kinetic + np.testing.assert_allclose(comps[2], 16.0 / 3.0) # ion thermal + np.testing.assert_allclose(comps[3], 2.0) # ion kinetic + np.testing.assert_allclose(comps[4], 0.5) # electric + np.testing.assert_allclose(comps[5], 2.0) # magnetic + np.testing.assert_allclose(comps[6], comps[:6].sum()) # total + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + elc = _make_5mom(rho=1.0, vx=0.0, p=1.0) + field = _make(_G1D, np.zeros((1, 6))) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + ms.energetics(d, elc, field) + + +class TestAccumulateCurrent: + + def _species(self): + return _make(_G1D, np.array([[1.0, 2.0, -3.0]])) + + def test_default_negates(self): + d = self._species() + out = ms.accumulate_current(d) + np.testing.assert_allclose(out.values, -d.values) + + def test_qbym_scales_by_charge_over_mass(self): + d = self._species() + out = ms.accumulate_current(d, qbym=True, charge=2.0, mass=4.0) + np.testing.assert_allclose(out.values, 0.5 * d.values) + + def test_qbym_negative_charge(self): + d = self._species() + out = ms.accumulate_current(d, qbym=True, charge=-1.0, mass=2.0) + np.testing.assert_allclose(out.values, -0.5 * d.values) + + def test_qbym_without_mass_raises(self): + d = self._species() + with pytest.raises(ValueError, match="qbym"): + ms.accumulate_current(d, qbym=True, charge=2.0) # mass missing + + def test_qbym_without_charge_raises(self): + d = self._species() + with pytest.raises(ValueError, match="qbym"): + ms.accumulate_current(d, qbym=True, mass=4.0) # charge missing + + def test_inplace_mutates(self): + d = self._species() + out = ms.accumulate_current(d, inplace=True) + assert out is d + + def test_grid_passed_through(self): + d = self._species() + out = ms.accumulate_current(d) + np.testing.assert_allclose(out.grid[0], _G1D[0]) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + ms.accumulate_current(d) + + +class TestAccumulateCurrentPrivateHelperFallback: + """``_accumulate_current`` (the moved array-level ``models.energetics`` + function) still silently falls back to the ``qbym=False`` formula when + ``mass``/``charge`` are missing -- the public verb now refuses that + combination before ever calling through (see ``accumulate_current``'s own + qbym guard above), so this behavior is only reachable by calling the + private helper directly, exactly as the pre-restructure + ``tests_models_energetics.py`` did against ``models.accumulate_current``.""" + + def test_qbym_without_mass_falls_back_to_negation(self): + values = np.array([[1.0, 2.0, 3.0]]) + _, out = ms._accumulate_current(_G1D, + values, + qbym=True, + charge=-1.0, + mass=None) + np.testing.assert_allclose(out, -values) + + def test_qbym_without_charge_falls_back_to_negation(self): + values = np.array([[1.0, 2.0, 3.0]]) + _, out = ms._accumulate_current(_G1D, + values, + qbym=True, + charge=None, + mass=1.0) + np.testing.assert_allclose(out, -values) diff --git a/tests/test_diagnostics_pkpm.py b/tests/test_diagnostics_pkpm.py new file mode 100644 index 00000000..f07d4ad7 --- /dev/null +++ b/tests/test_diagnostics_pkpm.py @@ -0,0 +1,246 @@ +"""Tests for postgkyl.diagnostics.pkpm -- PKPM Laguerre-moment composition, +folding the array-math analytic tests (formerly tests_models_laguerre.py) +with the verb-level guard/inplace tests (formerly part of +tests_ops_physics.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics import pkpm +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +def _square_inputs(n=5): + x = np.linspace(0.0, 1.0, n + 1) + vpar = np.linspace(-2.0, 2.0, n + 1) + f_values = np.ones((n, n, 2)) + t_over_m_values = np.ones((n, n, 1)) + return [x, vpar], f_values, t_over_m_values + + +class TestLaguerreComposePrivateHelperShape: + """Ported directly against the private array-level ``_laguerre_compose`` + (rather than the public ``GData``-facing verb), because these fixtures use + a T/m field that spans both ``x`` and ``vpar`` (``(n, n, 1)``, matching the + original ``tests_models_laguerre.py`` array-level fixture) -- physically + unrealistic for PKPM's actual T/m (a configuration-space-only quantity), + and it excites the broadcast bug (see ``pkpm``'s ``_laguerre_compose`` + docstring note) enough to make the returned array's spatial-axis count + (4) disagree with its own returned grid's length (3), which + ``GDataState.push``/``set_grid`` (correctly) refuses to accept. The + physically-sane T/m-on-``x``-only fixture used in the tests below (and in + ``TestLaguerreCompose``) does not hit this inconsistency; see there for the + public-verb-level tests.""" + + def test_output_grid_has_three_axes(self): + grid, f_values, t_m = _square_inputs() + out_grid, _ = pkpm._laguerre_compose(grid, f_values, t_m) + assert len(out_grid) == 3 + + def test_output_has_component_axis(self): + grid, f_values, t_m = _square_inputs() + _, out_f = pkpm._laguerre_compose(grid, f_values, t_m) + assert out_f.shape[-1] == 1 + + def test_third_axis_is_copy_of_vpar(self): + grid, f_values, t_m = _square_inputs() + out_grid, _ = pkpm._laguerre_compose(grid, f_values, t_m) + np.testing.assert_allclose(out_grid[2], grid[1]) + + def test_g_zero_reduces_to_maxwellian_of_f0(self): + # G = 0 -> F1 = F0, so f = F0*(2 - vperp^2/(2*T_m))/(2*pi*T_m) * + # exp(-vperp^2/(2*T_m)). + # + # T_m is broadcast against the 3-D (x, vpar, vperp) meshgrid with an + # extra np.newaxis (`T_m[..., np.newaxis, np.newaxis]`), one more than + # vperp_3D's single new axis -- inherited verbatim from + # src_bak/postgkyl/tools/laguerre_compose.py via + # postgkyl/diagnostics/pkpm's ``_laguerre_compose``, this makes the + # returned array 4 spatial axes deep (with a spurious, constant-along- + # itself extra axis) instead of the 3 the docstring/grid describe; the + # legacy test corpus never checked this middle shape either, only + # ``len(out_grid)`` and the trailing component axis, so this is a + # preexisting, untested quirk, not a regression -- reproduced here + # rather than silently corrected. + n = 4 + x = np.linspace(0.0, 1.0, n + 1) + vpar = np.linspace(-1.0, 1.0, n + 1) + F0_val, T_m_val = 2.0, 1.5 + f_values = np.zeros((n, n, 2)) + f_values[..., 0] = F0_val + t_over_m_values = np.full((n, n, 1), T_m_val) + + _, f = pkpm._laguerre_compose([x, vpar], f_values, t_over_m_values) + assert f.shape == (n, n, n, n, 1) + vperp_cc = 0.5 * (vpar[:-1] + vpar[1:]) + expected = (F0_val * (2 - vperp_cc**2 / (2 * T_m_val)) / + (2 * np.pi * T_m_val) * np.exp(-(vperp_cc**2) / (2 * T_m_val))) + # Every (x_cc, vpar_cc, spurious-axis) slice reproduces the same + # vperp-dependent curve. + np.testing.assert_allclose(f[0, 0, 0, :, 0], expected, rtol=1e-10) + np.testing.assert_allclose(f[0, 0, 2, :, 0], expected, rtol=1e-10) + + +class TestLaguerreCompose: + + def test_matches_private_helper(self): + x = np.linspace(0.0, 1.0, 3) # 2 cells + vpar = np.linspace(-1.0, 1.0, 3) # 2 cells + f_values = np.zeros((2, 2, 2)) + f_values[..., 0] = 1.0 # F0 + f_values[..., 1] = 0.5 # G + f = _make([x, vpar], f_values) + t_over_m = _make([x], np.full((2, 1), 2.0)) + + out = pkpm.laguerre_compose(f, t_over_m) + grid, values = pkpm._laguerre_compose(f.grid, f.values, t_over_m.values) + for d in range(len(grid)): + np.testing.assert_allclose(out.grid[d], grid[d]) + np.testing.assert_allclose(out.values, values) + + def test_extends_grid_with_vperp(self): + x = np.linspace(0.0, 1.0, 3) + vpar = np.linspace(-1.0, 1.0, 3) + f_values = np.zeros((2, 2, 2)) + f_values[..., 0] = 1.0 + f_values[..., 1] = 0.5 + f = _make([x, vpar], f_values) + t_over_m = _make([x], np.full((2, 1), 2.0)) + out = pkpm.laguerre_compose(f, t_over_m) + assert len(out.grid) == 3 + np.testing.assert_allclose(out.grid[2], + f.grid[1]) # vperp is a copy of vpar + + def test_inplace_mutates_distribution(self): + x = np.linspace(0.0, 1.0, 3) + vpar = np.linspace(-1.0, 1.0, 3) + f_values = np.zeros((2, 2, 2)) + f_values[..., 0] = 1.0 + f_values[..., 1] = 0.5 + f = _make([x, vpar], f_values) + t_over_m = _make([x], np.full((2, 1), 2.0)) + out = pkpm.laguerre_compose(f, t_over_m, inplace=True) + assert out is f + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + t_over_m = _make([np.array([0.0, 1.0])], np.array([[2.0]])) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pkpm.laguerre_compose(d, t_over_m) + + +# ---------------------------------------------------------------- load_pkpm +# No pkpm fixture is staged under tests/test_data, and postgkyl's own .gkyl +# writer (io/writer.py) does not reproduce the "multi-range" (file_type 3) +# structure the *compiled* reader (GkylCReader, tried first whenever the +# shim is available) expects for a real Gkeyll file -- so a naively +# write-then-load'ed synthetic file bounces off ``gpython_read_field`` before +# ``load_pkpm`` ever sees it. Following the same technique +# ``tests_bak/test_gk_load_quantity.py`` used for the (equally ctypes-only) +# old gk_quantities registry, the synthetic PKPM/pkpm_vars datasets are +# served in-memory by monkeypatching the ``GData`` name in ``pkpm`` itself +# calls -- this exercises the *real* naming convention, interpolation, +# ``laguerre_compose``, and ``transform_frame`` pipeline end to end; only +# the on-disk-file-format step is stubbed. +@needs_gkeyll +class TestLoadPkpm: + + _NB_HYBRID_2D_P1 = 6 # gpython.basis.num_basis("hybrid", 2, 1) + _NB_SER_1D_P1 = 2 # gpython.basis.num_basis("serendipity", 1, 1) + + def _synthetic_gf(self, F0=3.0, G=1.0): + """Two-field (F0, G) PKPM distribution on a 2-cell (x, vpar) grid; only + the mean coefficient is populated, per field, per dimension, so the + interpolated field value is exactly ``F0``/``G`` everywhere (the mean + basis function is ``2**(-ndim/2)``, so ``coeff0 = value * 2**(ndim/2)``).""" + nb = self._NB_HYBRID_2D_P1 + x = np.linspace(0.0, 1.0, 3) + vpar = np.linspace(-1.0, 1.0, 3) + values = np.zeros((2, 2, 2 * nb)) + values[..., 0 * nb] = F0 * 2.0**(2 / 2) + values[..., 1 * nb] = G * 2.0**(2 / 2) + g = pg.GData(ctx={"poly_order": 1, "basis_type": "hybrid"}) + g.push([x, vpar], values) + return g + + def _synthetic_gvars(self, u=(0.1, 0.2, 0.3), t_over_m=2.0): + """4-component (ux, uy, uz, T/m) PKPM variables on the same 1-D (x) grid.""" + nb = self._NB_SER_1D_P1 + x = np.linspace(0.0, 1.0, 3) + values = np.zeros((2, nb * 4)) + for i, uc in enumerate(u): + values[:, i * nb] = uc * 2.0**0.5 + values[:, 3 * nb] = t_over_m * 2.0**0.5 + g = pg.GData(ctx={"poly_order": 1, "basis_type": "serendipity"}) + g.push([x], values) + return g + + def _patch(self, monkeypatch, gf, gvars): + + def fake_ctor(file_name, **kwargs): + return gvars if "pkpm_vars" in file_name else gf + + monkeypatch.setattr(pkpm, "GData", fake_ctor) + + def test_output_grid_gains_vperp(self, monkeypatch): + gf, gvars = self._synthetic_gf(), self._synthetic_gvars() + self._patch(monkeypatch, gf, gvars) + out = pkpm.load_pkpm("sim", "ion", 0, 1) + # x, vpar, vperp: transform_frame shifts vpar/vperp per cell by the bulk + # velocity, so (unlike pre-transform) they are no longer identical, but + # both gained the same third (meshgrid) shape. + assert len(out.get_grid()) == 3 + assert out.get_grid()[1].shape == out.get_grid()[2].shape + + def test_matches_manual_compose_and_transform(self, monkeypatch): + F0, G, u, t_over_m = 3.0, 1.0, (0.1, 0.2, 0.3), 2.0 + gf, gvars = self._synthetic_gf(F0, G), self._synthetic_gvars(u, t_over_m) + self._patch(monkeypatch, gf, gvars) + out = pkpm.load_pkpm("sim", "ion", 0, 1) + + gf_interpolated = gf.interpolate() + gvars_interpolated = gvars.interpolate() + composed = pkpm.laguerre_compose(gf_interpolated, + gvars_interpolated.select(comp=3)) + from postgkyl.diagnostics.vm.kinetic import transform_frame + expected = transform_frame(composed, + gvars_interpolated.select(comp="0:3"), + cdim=1) + + np.testing.assert_allclose(out.values, expected.values) + for d in range(3): + np.testing.assert_allclose(out.get_grid()[d], expected.get_grid()[d]) + + def test_tag_and_label(self, monkeypatch): + gf, gvars = self._synthetic_gf(), self._synthetic_gvars() + self._patch(monkeypatch, gf, gvars) + out = pkpm.load_pkpm("sim", "ion", 0, 1, tag="mytag", label="mylabel") + assert out.get_tag() == "mytag" + assert out.get_label() == "mylabel" + + def test_default_tag_and_label(self, monkeypatch): + gf, gvars = self._synthetic_gf(), self._synthetic_gvars() + self._patch(monkeypatch, gf, gvars) + out = pkpm.load_pkpm("sim", "ion", 0, 1) + assert out.get_tag() == "default" diff --git a/tests/test_diagnostics_plasma.py b/tests/test_diagnostics_plasma.py new file mode 100644 index 00000000..6f8b6435 --- /dev/null +++ b/tests/test_diagnostics_plasma.py @@ -0,0 +1,284 @@ +"""Tests for postgkyl.diagnostics.mom.plasma -- plasma-parameter GData verbs +(magB, vt, vA, omegaC, omegaP, d, lambdaD, rho, beta), porting the analytic +array-math assertions of tests_models_plasma_params.py onto the new +GData-facing wrappers -- these functions never had a verb layer before this +restructure, so there is no old ops-level dispatch to preserve.""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest +import scipy.constants as const + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.mom import plasma as pp +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +_G1 = [np.array([0.0, 1.0])] + +# EM field: [Ex, Ey, Ez, Bx, By, Bz] Bx=3, By=4, Bz=0 -> |B|=5 +_FIELD_VALS = np.array([[0.0, 0.0, 0.0, 3.0, 4.0, 0.0]]) +_MAGB = 5.0 + +# 5-moment species: rho=2, vx=0.5, vy=0, vz=0, p=0.6 +_GAMMA = 5.0 / 3.0 +_RHO = 2.0 +_VX = 0.5 +_P = 0.6 +_E = _P / (_GAMMA - 1) + 0.5 * _RHO * _VX**2 +_MOM5 = np.array([[_RHO, _RHO * _VX, 0.0, 0.0, _E]]) + + +def _field(): + return _make(_G1, _FIELD_VALS) + + +def _species(): + return _make(_G1, _MOM5) + + +class TestMagB: + + def test_magnitude(self): + out = pp.magB(_field()) + np.testing.assert_allclose(out.values.flat[0], _MAGB, rtol=1e-10) + + def test_inplace_mutates_field(self): + field = _field() + out = pp.magB(field, inplace=True) + assert out is field + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.magB(d) + + +class TestVt: + + def test_no_sqrt2_defaults_false(self): + out = pp.vt(_species()) + T = _P / _RHO + np.testing.assert_allclose(out.values.flat[0], np.sqrt(2.0 * T), rtol=1e-10) + + def test_no_sqrt2(self): + out = pp.vt(_species(), no_sqrt2=True) + T = _P / _RHO + np.testing.assert_allclose(out.values.flat[0], np.sqrt(T), rtol=1e-10) + + def test_mass_scales_result(self): + out = pp.vt(_species(), mass=2.0, no_sqrt2=True) + T = _P / _RHO + np.testing.assert_allclose(out.values.flat[0], np.sqrt(T / 2.0), rtol=1e-10) + + def test_mhd_uses_mhd_temperature(self): + bx, by, bz = 1.0, 0.0, 0.0 + mag_p = 0.5 * (bx**2 + by**2 + bz**2) + e_mhd = 0.5 * _RHO * _VX**2 + _P / (_GAMMA - 1) + mag_p + mhd_vals = np.array([[_RHO, _RHO * _VX, 0.0, 0.0, e_mhd, bx, by, bz]]) + d = _make(_G1, mhd_vals) + out = pp.vt(d, gas_gamma=_GAMMA, mhd=True, no_sqrt2=True) + np.testing.assert_allclose(out.values.flat[0], + np.sqrt(_P / _RHO), + rtol=1e-10) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.vt(d) + + +class TestVA: + + def test_alfven_speed(self): + out = pp.vA(_species(), _field()) + expected = _MAGB / np.sqrt(_RHO) + np.testing.assert_allclose(out.values.flat[0], expected, rtol=1e-10) + + def test_mu0_scales_result(self): + out = pp.vA(_species(), _field(), mu_0=2.0) + expected = _MAGB / np.sqrt(2.0 * _RHO) + np.testing.assert_allclose(out.values.flat[0], expected, rtol=1e-10) + + def test_result_carries_species_grid(self): + species, field = _species(), _field() + out = pp.vA(species, field, inplace=True) + assert out is species + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + field = _field() + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.vA(d, field) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.vA(_species(), d) + + +class TestOmegaC: + + def test_cyclotron_frequency(self): + out = pp.omegaC(_field(), mass=1.0, charge=1.0) + np.testing.assert_allclose(out.values.flat[0], _MAGB, rtol=1e-10) + + def test_uses_absolute_charge(self): + oC_pos = pp.omegaC(_field(), mass=1.0, charge=1.0) + oC_neg = pp.omegaC(_field(), mass=1.0, charge=-1.0) + np.testing.assert_allclose(oC_pos.values.flat[0], + oC_neg.values.flat[0], + rtol=1e-10) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.omegaC(d) + + +class TestOmegaP: + + def test_plasma_frequency(self): + out = pp.omegaP(_species(), mass=1.0, charge=1.0, epsilon_0=1.0) + expected = np.sqrt(_RHO) + np.testing.assert_allclose(out.values.flat[0], expected, rtol=1e-10) + + def test_hydrogen_matches_nrl_formulary(self): + # NRL Plasma Formulary: f_pi[Hz] = 2.1e2 * Z * sqrt(n[cm^-3] / mu) for a + # singly-charged ion of mass number mu; compare our SI computation + # (mass density rho = n * m_p, as fluid moment data stores it) against + # this textbook approximation to its own (2-digit) precision. + n = 1.0e20 # m^-3 + rho_vals = np.array([[n * const.m_p]]) + d = _make(_G1, rho_vals) + out = pp.omegaP(d, + mass=const.m_p, + charge=const.e, + epsilon_0=const.epsilon_0) + expected_exact = np.sqrt(n * const.e**2 / (const.epsilon_0 * const.m_p)) + np.testing.assert_allclose(out.values.flat[0], expected_exact, rtol=1e-9) + + n_cm3 = n * 1e-6 + omega_nrl = 2 * np.pi * 2.1e2 * np.sqrt(n_cm3) + np.testing.assert_allclose(out.values.flat[0], omega_nrl, rtol=5e-3) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.omegaP(d) + + +class TestD: + + def test_skin_depth(self): + dd = pp.d(_species(), mass=1.0, charge=1.0, epsilon_0=1.0, mu_0=1.0) + omegaP = pp.omegaP(_species(), mass=1.0, charge=1.0, epsilon_0=1.0) + expected = 1.0 / omegaP.values.flat[0] + np.testing.assert_allclose(dd.values.flat[0], expected, rtol=1e-10) + + @needs_gkeyll + def test_rejects_modal_data(self): + modal = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.d(modal) + + +class TestLambdaD: + + def test_debye_length(self): + out = pp.lambdaD(_species(), + mass=1.0, + charge=1.0, + epsilon_0=1.0, + mu_0=1.0, + no_sqrt2=False) + vt_out = pp.vt(_species(), no_sqrt2=False) + omegaP_out = pp.omegaP(_species(), mass=1.0, charge=1.0, epsilon_0=1.0) + expected = vt_out.values.flat[0] / omegaP_out.values.flat[0] / np.sqrt(2.0) + np.testing.assert_allclose(out.values.flat[0], expected, rtol=1e-10) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.lambdaD(d) + + +class TestRho: + + def test_larmor_radius(self): + out = pp.rho(_species(), _field(), mass=1.0, charge=1.0, no_sqrt2=False) + vt_out = pp.vt(_species(), no_sqrt2=False) + omegaC_out = pp.omegaC(_field(), mass=1.0, charge=1.0) + expected = vt_out.values.flat[0] / omegaC_out.values.flat[0] + np.testing.assert_allclose(out.values.flat[0], expected, rtol=1e-10) + + def test_no_sqrt2_matches_default_after_normalization(self): + rho_default = pp.rho(_species(), _field(), mass=1.0, charge=1.0) + rho_no_sqrt2 = pp.rho(_species(), + _field(), + mass=1.0, + charge=1.0, + no_sqrt2=True) + np.testing.assert_allclose(rho_no_sqrt2.values.flat[0] / + rho_default.values.flat[0], + 1.0, + rtol=1e-8) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + field = _field() + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.rho(d, field) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.rho(_species(), d) + + +class TestBeta: + + def test_plasma_beta(self): + out = pp.beta(_species(), _field(), mu_0=1.0, no_sqrt2=False) + vt_out = pp.vt(_species(), no_sqrt2=False) + vA_out = pp.vA(_species(), _field(), mu_0=1.0) + expected = vt_out.values.flat[0]**2 / vA_out.values.flat[0]**2 + np.testing.assert_allclose(out.values.flat[0], expected, rtol=1e-10) + + def test_no_sqrt2_matches_default(self): + # The "* 2.0" correction for no_sqrt2=True exactly compensates for the + # missing sqrt(2) factor squared in v_th**2, so both conventions give + # the same beta. + beta_default = pp.beta(_species(), _field(), mu_0=1.0) + beta_no_sqrt2 = pp.beta(_species(), _field(), mu_0=1.0, no_sqrt2=True) + np.testing.assert_allclose(beta_no_sqrt2.values.flat[0], + beta_default.values.flat[0], + rtol=1e-10) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + field = _field() + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.beta(d, field) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + pp.beta(_species(), d) diff --git a/tests/test_diagnostics_programs_energy_balance.py b/tests/test_diagnostics_programs_energy_balance.py new file mode 100644 index 00000000..d13537fc --- /dev/null +++ b/tests/test_diagnostics_programs_energy_balance.py @@ -0,0 +1,355 @@ +"""Tests for ``postgkyl.diagnostics.gk.energy_balance``. + +Ported/extended from ``src_bak/postgkyl/apps/gk_energy_balance.py`` (no +``tests_bak`` corpus exists for this app -- it was never covered upstream). +The repo does not ship a multi-file gyrokinetic energy-balance fixture set +(``-field_energy_dot.gkyl``, ``..._fdot_integrated_moms.gkyl``, ...), +so the full figure path is exercised against synthetic per-file datasets +stubbed through ``utils.GData`` (the same technique +``tests/test_diagnostics_gk_load.py`` uses for the quantity registry), +rather than skipped outright -- this gives real coverage of the block/ +species accumulation loop and both (absolute- and relative-error) plotting +branches. The pure residual formula and accumulation helper are unit-tested +directly with no I/O at all. + +Run: PYTHONPATH=src pytest tests/test_diagnostics_programs_energy_balance.py -v +""" + +from __future__ import annotations + +import importlib +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from postgkyl.diagnostics.gk import utils as gk_utils + +eb = importlib.import_module("postgkyl.diagnostics.gk.energy_balance") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") + + +class _FakeGData: + """Stands in for ``postgkyl.gdata.GData`` -- just enough surface for + ``utils.read_gfile``/``read_gfile_if_present`` (``get_grid``/``get_values``/ + ``ctx``).""" + + def __init__(self, grid, values, ctx=None): + self._grid = grid + self._values = values + self.ctx = ctx or {} + + def get_grid(self): + return self._grid + + def get_values(self): + return self._values + + +class _StubFiles: + """Registers ``(grid, values)`` for a set of file names and monkeypatches + ``utils.GData`` to serve them, touching each file on disk so the + existence checks in ``read_gfile_if_present`` pass.""" + + def __init__(self, tmp_path, monkeypatch): + self._tmp_path = tmp_path + self._registry: dict[str, _FakeGData] = {} + monkeypatch.setattr(gk_utils, "GData", self._dispatch) + + def _dispatch(self, file_name): + return self._registry[file_name] + + def add(self, file_name: str, time_edges: np.ndarray, + values: np.ndarray) -> None: + open(file_name, "w").close() + self._registry[file_name] = _FakeGData([np.asarray(time_edges)], + np.asarray(values)) + + +@pytest.fixture +def stub(tmp_path, monkeypatch): + return _StubFiles(tmp_path, monkeypatch) + + +def _build_sim(stub, + tmp_path, + name="sim", + species=("ion", ), + *, + with_src=True, + with_bflux=True, + with_apar=False, + n=5): + """Populate a minimal single-block energy-balance file set.""" + path = str(tmp_path) + "/" + # A dynvector's grid is exactly one time stamp per recorded sample (see + # ``io/gkyl_reader.py``'s ``_read_t2_v1``), not N+1 cell edges like a + # field file -- the fake GData below mimics that real convention. + time = np.linspace(0.0, 1.0, n) + + for sp in species: + fdot_vals = np.zeros((n, 3)) + fdot_vals[:, 2] = np.linspace(1.0, 2.0, n) + stub.add(f"{path}{name}-{sp}_fdot_integrated_moms.gkyl", time, fdot_vals) + + if with_src: + src_vals = np.zeros((n, 3)) + src_vals[:, 2] = 0.1 + stub.add(f"{path}{name}-{sp}_source_integrated_moms.gkyl", time, src_vals) + if with_bflux: + bflux_vals = np.zeros((n, 3)) + bflux_vals[:, 2] = 0.05 + stub.add( + f"{path}{name}-{sp}_bflux_xlower_integrated_HamiltonianMoments.gkyl", + time, bflux_vals) + + field_dot_vals = np.zeros((n, 1)) + field_dot_vals[:, 0] = 0.2 + stub.add(f"{path}{name}-field_energy_dot.gkyl", time, field_dot_vals) + + if with_apar: + apar_dot_vals = np.zeros((n, 1)) + apar_dot_vals[:, 0] = 0.15 + stub.add(f"{path}{name}-apar_energy_dot.gkyl", time, apar_dot_vals) + + return path + + +class TestEnergyBalanceErrorPure: + """The residual formula -- pure array arithmetic, no I/O.""" + + def test_no_apar(self): + fdot = np.array([2.0, 3.0]) + src = np.array([1.0, 1.0]) + bflux = np.array([0.5, 0.5]) + field_dot = np.array([1.0, 1.0]) + err = eb.energy_balance_error(fdot, src, bflux, field_dot) + np.testing.assert_allclose(err, src - bflux - (fdot - field_dot)) + + def test_with_apar(self): + fdot = np.array([2.0]) + src = np.array([1.0]) + bflux = np.array([0.5]) + field_dot = np.array([1.0]) + apar_dot = np.array([0.25]) + err = eb.energy_balance_error(fdot, src, bflux, field_dot, apar_dot) + np.testing.assert_allclose(err, src - bflux - (fdot - field_dot - apar_dot)) + + +class TestAccumulatePure: + + def test_first_use_copies_not_aliases(self): + a = np.array([1.0, 2.0]) + out = eb._accumulate(None, a) + out[0] = 99.0 + assert a[0] == 1.0 + + def test_accumulates_sum(self): + out = eb._accumulate(np.array([1.0, 2.0]), np.array([3.0, 4.0])) + np.testing.assert_allclose(out, [4.0, 6.0]) + + +class TestResolvePure: + + def test_no_override_uses_default(self): + assert eb._resolve("/p/", None, "default.gkyl", 0) == "default.gkyl" + + def test_override_substitutes_block_then_species(self): + out = eb._resolve("/p/", "custom_*_*.gkyl", "unused", 3, "ion") + assert out == "/p/custom_3_ion.gkyl" + + +class TestGkEnergyBalanceSynthetic: + """Full figure-path coverage against stubbed per-file datasets.""" + + def test_full_path_with_src_and_bflux(self, stub, tmp_path): + path = _build_sim(stub, tmp_path) + fig, traces = eb.energy_balance("sim", ["ion"], path=path) + try: + assert traces.src is not None + assert traces.bflux_tot is not None + assert traces.mom_err is not None + assert traces.time.shape[0] == 5 + # src[0] is zeroed before computing the residual. + assert traces.mom_err.shape == (5, ) + finally: + plt.close(fig) + + def test_missing_source_and_bflux(self, stub, tmp_path): + path = _build_sim(stub, tmp_path, with_src=False, with_bflux=False) + fig, traces = eb.energy_balance("sim", ["ion"], path=path) + try: + assert traces.src is None + assert traces.bflux_tot is None + finally: + plt.close(fig) + + def test_electromagnetic_branch(self, stub, tmp_path): + path = _build_sim(stub, tmp_path, with_apar=True) + fig, traces = eb.energy_balance("sim", ["ion"], path=path) + try: + assert traces.apar_dot is not None + finally: + plt.close(fig) + + def test_multi_species_sums(self, stub, tmp_path): + path = _build_sim(stub, tmp_path, species=("ion", "elc")) + fig, traces = eb.energy_balance("sim", ["ion", "elc"], path=path) + try: + # Two identical species contributions sum to double a single one. + single_dir = tmp_path / "single" + single_dir.mkdir() + single_path = _build_sim(stub, single_dir, species=("ion", )) + _, single_traces = eb.energy_balance("sim", ["ion"], path=single_path) + np.testing.assert_allclose(traces.fdot, 2 * single_traces.fdot) + finally: + plt.close(fig) + + def test_relative_error_branch(self, stub, tmp_path): + path = _build_sim(stub, tmp_path) + n = 5 + time = np.linspace(0.0, 1.0, n) + field_vals = np.full((n, 1), 3.0) + stub.add(f"{path}sim-field_energy.gkyl", time, field_vals) + f_vals = np.zeros((n, 3)) + f_vals[:, 2] = 10.0 + stub.add(f"{path}sim-ion_integrated_moms.gkyl", time, f_vals) + # dt.gkyl records the timestep *between* frames, so it naturally has one + # fewer entry than the per-frame traces -- matching src_bak, which slices + # every per-frame trace with [1:] but never slices dt itself. + dt_time = np.linspace(0.0, 1.0, n - 1) + dt_vals = np.full((n - 1, 1), 0.2) + stub.add(f"{path}sim-dt.gkyl", dt_time, dt_vals) + + fig, traces = eb.energy_balance("sim", ["ion"], + path=path, + relative_error=True) + try: + assert traces.mom_err is None + assert traces.mom_err_norm is not None + # One point is dropped (t=0) relative to the absolute-error path. + assert traces.mom_err_norm.shape[0] == n - 1 + finally: + plt.close(fig) + + def test_relative_error_electromagnetic_absy_and_saveas(self, stub, tmp_path): + """Covers the apar branch inside the relative-error path together with + ``absy``/``saveas``/``show``.""" + path = _build_sim(stub, tmp_path, with_apar=True) + n = 5 + time = np.linspace(0.0, 1.0, n) + stub.add(f"{path}sim-field_energy.gkyl", time, np.full((n, 1), 3.0)) + stub.add(f"{path}sim-apar_energy.gkyl", time, np.full((n, 1), 1.0)) + f_vals = np.zeros((n, 3)) + f_vals[:, 2] = 10.0 + stub.add(f"{path}sim-ion_integrated_moms.gkyl", time, f_vals) + dt_time = np.linspace(0.0, 1.0, n - 1) + stub.add(f"{path}sim-dt.gkyl", dt_time, np.full((n - 1, 1), 0.2)) + + out_path = str(tmp_path / "out.png") + fig, traces = eb.energy_balance("sim", ["ion"], + path=path, + relative_error=True, + absy=True, + logy=True, + saveas=out_path) + try: + assert traces.mom_err_norm is not None + assert os.path.exists(out_path) + finally: + plt.close(fig) + + def test_relative_error_apar_dot_present_without_apar_energy( + self, stub, tmp_path): + """Regression for C1: a run can ship ``apar_energy_dot.gkyl`` (read in + the unrelated, earlier per-block loop that sets ``has_apar_dot``) + without shipping ``apar_energy.gkyl`` (read inside the relative-error + branch's own loop, which sets ``has_apar``). The relative-error branch + must gate every apar-dependent line -- the ``[1:]`` slicing, the + ``energy_balance_error`` call, and the ``denom`` computation -- on + ``has_apar``, not ``has_apar_dot``; gating on the wrong flag leaves + ``apar`` as ``None`` (never accumulated, since ``has_apar`` is False) + while still trying to slice it, raising + ``TypeError: 'NoneType' object is not subscriptable``.""" + path = _build_sim(stub, tmp_path, + with_apar=True) # stages apar_energy_dot.gkyl only. + n = 5 + time = np.linspace(0.0, 1.0, n) + # No "sim-apar_energy.gkyl" staged -- has_apar stays False. + stub.add(f"{path}sim-field_energy.gkyl", time, np.full((n, 1), 3.0)) + f_vals = np.zeros((n, 3)) + f_vals[:, 2] = 10.0 + stub.add(f"{path}sim-ion_integrated_moms.gkyl", time, f_vals) + dt_time = np.linspace(0.0, 1.0, n - 1) + stub.add(f"{path}sim-dt.gkyl", dt_time, np.full((n - 1, 1), 0.2)) + + fig, traces = eb.energy_balance("sim", ["ion"], + path=path, + relative_error=True) + try: + # No TypeError, and the electromagnetic term is correctly excluded + # (has_apar-gated) -- matches the electrostatic relative-error formula. + assert traces.mom_err_norm is not None + assert traces.mom_err_norm.shape[0] == n - 1 + finally: + plt.close(fig) + + def test_missing_required_field_dot_file_raises(self, stub, tmp_path): + path = str(tmp_path) + "/" + with pytest.raises(FileNotFoundError, match="field_energy_dot"): + eb.energy_balance("sim", ["ion"], path=path) + + def test_missing_required_fdot_file_raises(self, stub, tmp_path): + path = str(tmp_path) + "/" + n = 5 + time = np.linspace(0.0, 1.0, n) + stub.add(f"{path}sim-field_energy_dot.gkyl", time, np.zeros((n, 1))) + with pytest.raises(FileNotFoundError, match="fdot_integrated_moms"): + eb.energy_balance("sim", ["ion"], path=path) + + def test_bflux_override_and_absy_logy(self, stub, tmp_path): + path = _build_sim(stub, tmp_path, with_bflux=False) + n = 5 + time = np.linspace(0.0, 1.0, n) + override_vals = np.zeros((n, 3)) + override_vals[:, 2] = -0.05 + override_name = f"{path}custom_bflux.gkyl" + stub.add(override_name, time, override_vals) + + fig, traces = eb.energy_balance("sim", ["ion"], + path=path, + bflux_files={"xlower": "custom_bflux.gkyl"}, + absy=True, + logy=True, + show=True) + try: + assert traces.bflux_tot is not None + np.testing.assert_allclose(traces.bflux_tot, -0.05) + finally: + plt.close(fig) + + +class TestGkEnergyBalanceRealFixtures: + """Real end-to-end run against ``tests/test_data`` -- skipped loudly since + the repo does not ship a gyrokinetic energy-balance file family (only + single-frame distribution/geometry fixtures for a different diagnostic are + staged there).""" + + def test_real_fixture_energy_balance(self): + required = ("field_energy_dot.gkyl", "_fdot_integrated_moms.gkyl") + if not any( + any(f.endswith(suffix) for f in os.listdir(DATA)) + for suffix in required): + pytest.skip( + "tests/test_data ships no gyrokinetic energy-balance file family " + "(needs e.g. '-field_energy_dot.gkyl', " + "'-_fdot_integrated_moms.gkyl'); see " + "TestGkEnergyBalanceSynthetic for full-path coverage against " + "stubbed data instead.") + pytest.fail("fixture files appeared -- wire up a real-data assertion here") diff --git a/tests/test_diagnostics_programs_enstrophy.py b/tests/test_diagnostics_programs_enstrophy.py new file mode 100644 index 00000000..ad154eb7 --- /dev/null +++ b/tests/test_diagnostics_programs_enstrophy.py @@ -0,0 +1,123 @@ +"""Tests for ``postgkyl.diagnostics.mom.enstrophy``. + +Ported from ``src_bak/postgkyl/tools/calc_enstrophy.py`` (no ``tests_bak`` +corpus exists for this tool). The pure per-frame math (``_enstrophy_terms``) +is checked against an analytic velocity field where the curl and the +velocity-gradient tensor are hand-computable exactly (linear-in-coordinate +components, so ``np.gradient(..., edge_order=2)`` on a uniform grid +reproduces the analytic derivative exactly); the frame-sweep wiring +(``enstrophy``) is exercised against a synthetic multi-frame file family +stubbed through ``postgkyl.diagnostics.mom.enstrophy.GData`` -- the repo ships +no multi-frame 3-D five-moment ``.gkyl`` fixture family for this tool, so no +real-fixture path is attempted (see this layer's report). + +Run: PYTHONPATH=src pytest tests/test_diagnostics_programs_enstrophy.py -v +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.diagnostics.mom import enstrophy as ens + + +class _FakeGData: + + def __init__(self, grid, values): + self.grid = grid + self.values = values + + +class TestEnstrophyTermsAnalytic: + """u = x, v = y, w = -2z (irrotational, incompressible): curl is exactly + zero everywhere, and the velocity-gradient tensor's diagonal is constant + (1, 1, -2) everywhere, so both integrals are exactly computable by hand.""" + + def _field(self, n=4, rho0=2.0): + dx = dy = dz = 1.0 + coords = np.arange(n, dtype=np.float64) + x, y, z = np.meshgrid(coords, coords, coords, indexing="ij") + rho = np.full((n, n, n), rho0) + u, v, w = x, y, -2.0 * z + px, py, pz = u * rho, v * rho, w * rho + return rho, px, py, pz, dx, dy, dz + + def test_curl_is_zero_for_irrotational_field(self): + rho, px, py, pz, dx, dy, dz = self._field() + enstrophy_val, _ = ens._enstrophy_terms(rho, px, py, pz, dx, dy, dz) + np.testing.assert_allclose(enstrophy_val, 0.0, atol=1e-10) + + def test_incompressible_term_matches_hand_derivation(self): + n = 4 + rho, px, py, pz, dx, dy, dz = self._field(n=n, rho0=2.0) + _, incompressible = ens._enstrophy_terms(rho, px, py, pz, dx, dy, dz) + # diag(grad) = (1, 1, -2) everywhere -> trace(M^T (*) M) = 1^2+1^2+(-2)^2 = 6. + # incom_mag = 6 * rho = 12, summed only over the (n-1)^3 sub-cube the + # nested loop's `range(n - 1)` bound reaches (a quirk preserved verbatim + # from src_bak -- see the module docstring), times dx*dy*dz = 1. + expected = 6.0 * 2.0 * (n - 1)**3 + np.testing.assert_allclose(incompressible, expected) + + def test_zero_velocity_gives_zero_both_terms(self): + n = 3 + rho = np.full((n, n, n), 1.0) + zero = np.zeros((n, n, n)) + enstrophy_val, incompressible = ens._enstrophy_terms( + rho, zero, zero, zero, 1.0, 1.0, 1.0) + np.testing.assert_allclose(enstrophy_val, 0.0) + np.testing.assert_allclose(incompressible, 0.0) + + +class TestEnstrophySweep: + """Frame-sweep wiring: ``enstrophy()`` reads ``stem{frame}.ext`` for each + frame in ``[init_frame, final_frame]`` and stacks the per-frame results.""" + + def test_sweeps_expected_frame_range(self, monkeypatch): + n = 3 + edges = np.arange(n + 1, dtype=np.float64) + rho = np.full((n, n, n), 1.0) + + calls = [] + + def fake_gdata(file_name): + calls.append(file_name) + values = np.stack([rho, rho, rho, rho], axis=-1) # rho, px=py=pz=rho + return _FakeGData([edges, edges, edges], values) + + monkeypatch.setattr(ens, "GData", fake_gdata) + + out = ens.enstrophy("sim-fluid_", 2, 4, extension="dat") + # The first frame is read twice: once up front for the grid spacing, + # then again inside the sweep loop. + assert calls == [ + "sim-fluid_2.dat", "sim-fluid_2.dat", "sim-fluid_3.dat", + "sim-fluid_4.dat" + ] + assert out.enstrophy.shape == (3, ) + assert out.incompressible_enstrophy.shape == (3, ) + # u = v = w = px/rho = 1 (constant) -> zero curl and zero gradient. + np.testing.assert_allclose(out.enstrophy, 0.0) + np.testing.assert_allclose(out.incompressible_enstrophy, 0.0) + + def test_single_frame_range(self, monkeypatch): + n = 3 + edges = np.arange(n + 1, dtype=np.float64) + rho = np.full((n, n, n), 1.0) + + def fake_gdata(file_name): + values = np.stack([rho, rho, rho, rho], axis=-1) + return _FakeGData([edges, edges, edges], values) + + monkeypatch.setattr(ens, "GData", fake_gdata) + out = ens.enstrophy("sim-fluid_", 0, 0) + assert out.enstrophy.shape == (1, ) + + +class TestEnstrophyTracesIsFrozen: + + def test_fields_present(self): + t = ens.EnstrophyTraces(enstrophy=np.array([1.0]), + incompressible_enstrophy=np.array([2.0])) + with pytest.raises(Exception): + t.enstrophy = np.array([3.0]) diff --git a/tests/test_diagnostics_programs_ke_dke.py b/tests/test_diagnostics_programs_ke_dke.py new file mode 100644 index 00000000..99917361 --- /dev/null +++ b/tests/test_diagnostics_programs_ke_dke.py @@ -0,0 +1,145 @@ +"""Tests for ``postgkyl.diagnostics.mom.ke_dke``. + +Ported from ``src_bak/postgkyl/tools/calc_ke_dke.py`` (no ``tests_bak`` +corpus exists for this tool). See the module docstring for the three +``src_bak`` bugs this port fixes (a file-name f-string missing its own +parameter, an array-aliasing bug, and an off-by-one difference-loop bound) +-- the tests here pin the *fixed* behavior: an exact analytic kinetic-energy +value per frame, and a dissipation rate covering every consecutive frame +pair. + +Run: PYTHONPATH=src pytest tests/test_diagnostics_programs_ke_dke.py -v +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.diagnostics.mom import ke_dke as kd + + +class _FakeGData: + + def __init__(self, grid, values): + self.grid = grid + self.values = values + + +class TestKineticEnergyAnalytic: + + def test_uniform_velocity_matches_hand_derivation(self): + n = 4 + rho = np.full((n, n, n), 2.0) + u = np.full((n, n, n), 1.0) + v = np.full((n, n, n), 2.0) + w = np.full((n, n, n), 3.0) + px, py, pz = u * rho, v * rho, w * rho + dx = dy = dz = 0.5 + vol = 10.0 + ke = kd._kinetic_energy(rho, px, py, pz, dx, dy, dz, vol) + # e = rho*(u^2+v^2+w^2) = 2*(1+4+9) = 28 per cell, n^3 = 64 cells. + expected = 28.0 * (n**3) * dx * dy * dz * vol + np.testing.assert_allclose(ke, expected) + + def test_zero_velocity_gives_zero_energy(self): + n = 3 + rho = np.full((n, n, n), 5.0) + zero = np.zeros((n, n, n)) + ke = kd._kinetic_energy(rho, zero, zero, zero, 1.0, 1.0, 1.0, 1.0) + np.testing.assert_allclose(ke, 0.0) + + +class TestDissipationRatePure: + + def test_backward_difference_every_pair(self): + ke = np.array([1.0, 3.0, 6.0, 10.0]) + dke = kd._dissipation_rate(ke, dt=0.5) + expected = -(ke[1:] - ke[:-1]) / 0.5 + np.testing.assert_allclose(dke, expected) + assert dke.shape[0] == ke.shape[0] - 1 + + def test_constant_ke_gives_zero_dissipation(self): + ke = np.full(5, 3.0) + dke = kd._dissipation_rate(ke, dt=1.0) + np.testing.assert_allclose(dke, 0.0) + + +class TestKeDkeSweep: + + def _uniform_frame(self, n=3, value=1.0): + edges = np.arange(n + 1, dtype=np.float64) + rho = np.full((n, n, n), value) + values = np.stack([rho, rho, rho, rho], axis=-1) + return _FakeGData([edges, edges, edges], values) + + def test_sweeps_expected_frame_count_and_dke_length(self, monkeypatch): + calls = [] + + def fake_gdata(file_name): + calls.append(file_name) + return self._uniform_frame() + + monkeypatch.setattr(kd, "GData", fake_gdata) + + out = kd.ke_dke("sim-fluid_", + 0, + 3, + dim=3, + vol=1.0, + init_time=0.0, + final_time=3.0) + # First frame read twice (once for grid spacing, once in the sweep). + assert calls == [ + "sim-fluid_0.gkyl", "sim-fluid_0.gkyl", "sim-fluid_1.gkyl", + "sim-fluid_2.gkyl", "sim-fluid_3.gkyl" + ] + assert out.ke.shape == (4, ) + assert out.dke.shape == (3, ) + # u=v=w=1 (rho=1, px=py=pz=1/rho=... wait: px=py=pz=rho=1 -> u=v=w=1) + # constant across every frame -> dke is exactly zero, not just close. + np.testing.assert_allclose(out.dke, 0.0) + + def test_dim_2_uses_unit_z_spacing(self, monkeypatch): + + def fake_gdata(file_name): + return self._uniform_frame(n=2) + + monkeypatch.setattr(kd, "GData", fake_gdata) + out = kd.ke_dke("sim-fluid_", + 0, + 1, + dim=2, + vol=1.0, + init_time=0.0, + final_time=1.0) + assert out.ke.shape == (2, ) + + def test_uses_own_root_file_name_not_a_literal_string(self, monkeypatch): + """Regression test for the src_bak bug where the per-frame file name was + built as f"root_file_name{c:d}.gkyl" -- a literal string containing the + parameter's *name* -- instead of interpolating its value.""" + calls = [] + + def fake_gdata(file_name): + calls.append(file_name) + return self._uniform_frame() + + monkeypatch.setattr(kd, "GData", fake_gdata) + kd.ke_dke("distinctive_stem_", + 0, + 1, + dim=3, + vol=1.0, + init_time=0.0, + final_time=1.0) + assert all(c.startswith("distinctive_stem_") for c in calls) + assert not any("root_file_name" in c for c in calls) + + +class TestKineticEnergyTracesIsFrozen: + + def test_fields_present(self): + t = kd.KineticEnergyTraces(ke=np.array([1.0]), dke=np.array([])) + with pytest.raises(Exception): + t.ke = np.array([2.0]) diff --git a/tests/test_diagnostics_programs_nodes.py b/tests/test_diagnostics_programs_nodes.py new file mode 100644 index 00000000..741563a4 --- /dev/null +++ b/tests/test_diagnostics_programs_nodes.py @@ -0,0 +1,353 @@ +"""Tests for ``postgkyl.diagnostics.gk.nodes``. + +Ported from ``src_bak/postgkyl/apps/gk_nodes.py`` (no ``tests_bak`` corpus +exists for this app). The pure geometry helpers (``is_geo_mapc2p``, +``nodes_to_RZ``, ``multib_tag``, ``_parse_levels``) are unit-tested +unconditionally; the node-plotting figure path is exercised against +synthetic node arrays stubbed through ``utils.GData`` (single- and +multi-block). The poloidal-flux (``psi_file``) and wall overlays additionally +call ``GData.interpolate()`` on a *real* modal DG field (``nodes`` hardcodes +``poly_order=2``/basis ``"mt"`` for the psi read) -- the repo ships no +interpolatable p2-tensor poloidal-flux fixture, so that branch is skipped +loudly rather than faked. + +Run: PYTHONPATH=src pytest tests/test_diagnostics_programs_nodes.py -v +""" + +from __future__ import annotations + +import importlib +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from postgkyl.diagnostics.gk import utils as gk_utils + +nodes = importlib.import_module("postgkyl.diagnostics.gk.nodes") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GENERATED = os.path.join(DATA, "generated") + + +class TestGeometryEnum: + + def test_mapc2p_index_matches_gkeyll_header(self): + # gkeyll/core/zero/gkyl_eqn_type.h: GKYL_GEOMETRY_MAPC2P = 3. + assert nodes.GKYL_GEOMETRY_ID.index("GKYL_GEOMETRY_MAPC2P") == 3 + + +class TestIsGeoMapc2p: + + def test_defaults_true_when_absent(self): + assert nodes.is_geo_mapc2p({}) is True + + def test_true_for_mapc2p(self): + assert nodes.is_geo_mapc2p({"geometry_type": 3}) is True + + def test_false_for_tokamak(self): + assert nodes.is_geo_mapc2p({"geometry_type": 1}) is False + + +class TestNodesToRZ: + + def test_mapc2p_2d(self): + # A 3x2 grid of Cartesian (X, Y, Z) nodes on the unit circle at Z=0. + shape = (3, 2) + nodes_arr = np.zeros(shape + (3, )) + nodes_arr[..., 0] = 1.0 # X + nodes_arr[..., 1] = 0.0 # Y + nodes_arr[..., 2] = 5.0 # Z + major_r, vert_z = nodes.nodes_to_RZ(nodes_arr, is_mapc2p=True) + np.testing.assert_allclose(major_r, 1.0) + np.testing.assert_allclose(vert_z, 5.0) + + def test_non_mapc2p_2d(self): + shape = (3, 2) + nodes_arr = np.zeros(shape + (2, )) + nodes_arr[..., 0] = 2.0 # R + nodes_arr[..., 1] = -1.0 # Z + major_r, vert_z = nodes.nodes_to_RZ(nodes_arr, is_mapc2p=False) + np.testing.assert_allclose(major_r, 2.0) + np.testing.assert_allclose(vert_z, -1.0) + + def test_mapc2p_1d(self): + shape = (4, ) + nodes_arr = np.zeros(shape + (3, )) + nodes_arr[..., 0] = 3.0 + nodes_arr[..., 1] = 4.0 + nodes_arr[..., 2] = 7.0 + major_r, vert_z = nodes.nodes_to_RZ(nodes_arr, is_mapc2p=True) + np.testing.assert_allclose(major_r, 5.0) # sqrt(3^2+4^2) + np.testing.assert_allclose(vert_z, 7.0) + + def test_mapc2p_3d_slices_at_yidx_zero(self): + # cdim == 3 slices the y axis at index 0 before extracting X, Y, Z. + shape = (2, 3, 2) + nodes_arr = np.zeros(shape + (3, )) + nodes_arr[:, 0, :, 0] = 1.0 # X at y-index 0 + nodes_arr[:, 0, :, 1] = 0.0 # Y at y-index 0 + nodes_arr[:, 0, :, 2] = 9.0 # Z at y-index 0 + nodes_arr[:, 1, :, 0] = 100.0 # far-away values at y-index 1 (unused) + major_r, vert_z = nodes.nodes_to_RZ(nodes_arr, is_mapc2p=True) + np.testing.assert_allclose(major_r, 1.0) + np.testing.assert_allclose(vert_z, 9.0) + + +class TestMultibTag: + + def test_single_block_no_suffix(self): + assert nodes.multib_tag("nodes", 0, 1) == "nodes" + + def test_multiblock_suffix(self): + assert nodes.multib_tag("nodes", 2, 3) == "nodes_b2" + + +class TestParseLevels: + + def test_none_returns_cnlevels(self): + assert nodes._parse_levels(None, 11) == 11 + + def test_range_string(self): + out = nodes._parse_levels("0:1:3", 11) + np.testing.assert_allclose(out, [0.0, 0.5, 1.0]) + + def test_comma_list(self): + out = nodes._parse_levels("0.1,0.2,0.3", 11) + np.testing.assert_allclose(out, [0.1, 0.2, 0.3]) + + +class _FakeGData: + + def __init__(self, grid, values, ctx=None): + self._grid = grid + self._values = values + self.ctx = ctx or {} + + def get_grid(self): + return self._grid + + def get_values(self): + return self._values + + def interpolate(self, num_interp=None): + return self + + +class _StubFiles: + + def __init__(self, monkeypatch): + self._registry: dict[str, _FakeGData] = {} + monkeypatch.setattr(gk_utils, "GData", self._dispatch) + + def _dispatch(self, file_name, **kwargs): + return self._registry[file_name] + + def add(self, file_name: str, grid, values, ctx=None) -> None: + open(file_name, "w").close() + self._registry[file_name] = _FakeGData(grid, values, ctx) + + +@pytest.fixture +def stub(monkeypatch): + return _StubFiles(monkeypatch) + + +def _square_nodes(nx=3, ny=3): + """A simple 2-D mapc2p node grid: a regular (nx, ny) square in the X-Y + plane, with Z varying along x (so both the R and Z extents -- and hence + ``nodes``'s figure aspect ratio -- are nonzero and finite).""" + x = np.linspace(1.0, 2.0, nx) + y = np.linspace(0.0, 1.0, ny) + xx, yy = np.meshgrid(x, y, indexing="ij") + out = np.zeros((nx, ny, 3)) + out[..., 0] = xx + out[..., 1] = yy + out[..., 2] = xx # Z varies with x, giving a nonzero vertical extent. + return out + + +class TestGkNodesSynthetic: + + def test_single_block_no_overlays(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", + [np.arange(3.0), np.arange(3.0)], _square_nodes()) + fig = nodes.nodes("sim", path=path) + try: + assert fig is not None + assert len(fig.axes) == 1 + finally: + plt.close(fig) + + def test_multiblock_sums_extrema_across_blocks(self, stub, tmp_path): + path = str(tmp_path) + "/" + block0 = _square_nodes() + block1 = _square_nodes() + 5.0 # shifted far away in R and Z + stub.add(f"{path}sim_b0-nodes.gkyl", [np.arange(3.0)] * 2, block0) + stub.add(f"{path}sim_b1-nodes.gkyl", [np.arange(3.0)] * 2, block1) + fig = nodes.nodes("sim", path=path, multib="0,1") + try: + assert fig is not None + finally: + plt.close(fig) + + def test_non_mapc2p_geometry_type(self, stub, tmp_path): + path = str(tmp_path) + "/" + rz_nodes = np.zeros((3, 3, 2)) + rz_nodes[..., 0] = np.linspace(1.0, 2.0, 3)[:, None] + rz_nodes[..., 1] = np.linspace(-1.0, 1.0, 3)[None, :] + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, + rz_nodes, + ctx={"geometry_type": 1}) + fig = nodes.nodes("sim", path=path) + try: + assert fig is not None + finally: + plt.close(fig) + + def test_wall_file_overlay(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + wall_path = tmp_path / "wall.csv" + wall_path.write_text("0.0,0.0\n1.0,1.0\n2.0,0.0\n") + fig = nodes.nodes("sim", path=path, wall_file="wall.csv") + try: + assert fig is not None + finally: + plt.close(fig) + + def test_absolute_nodes_file_override(self, stub, tmp_path): + path = str(tmp_path) + "/" + abs_file = f"{path}custom_nodes.gkyl" + stub.add(abs_file, [np.arange(3.0)] * 2, _square_nodes()) + fig = nodes.nodes("sim", path=path, nodes_file=abs_file) + try: + assert fig is not None + finally: + plt.close(fig) + + def test_xlim_ylim_and_saveas(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + out_path = str(tmp_path / "out.png") + fig = nodes.nodes("sim", + path=path, + xlim=(0.0, 2.0), + ylim=(-1.0, 1.0), + saveas=out_path) + try: + assert fig.axes[0].get_xlim() == (0.0, 2.0) + assert fig.axes[0].get_ylim() == (-1.0, 1.0) + assert os.path.exists(out_path) + finally: + plt.close(fig) + + def test_1d_node_array_uses_line_plot_branch(self, stub, tmp_path): + path = str(tmp_path) + "/" + nodes_1d = np.zeros((4, 3)) + nodes_1d[:, 0] = np.linspace(1.0, 2.0, 4) + nodes_1d[:, 2] = np.linspace(0.0, 1.0, 4) + stub.add(f"{path}sim-nodes.gkyl", [np.arange(4.0)], nodes_1d) + fig = nodes.nodes("sim", path=path) + try: + assert fig is not None + finally: + plt.close(fig) + + def test_show_calls_plt_show(self, stub, tmp_path, monkeypatch): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + calls = [] + monkeypatch.setattr(plt, "show", lambda: calls.append(True)) + fig = nodes.nodes("sim", path=path, show=True) + try: + assert calls == [True] + finally: + plt.close(fig) + + +class TestGkNodesPsiOverlay: + """The ``psi_file`` overlay path, stubbed through ``gk_utils.GData`` (its + ``.interpolate()`` returns itself, carrying the already pre-shaped + edge-grid/cell-centered-values pair a caller registered) so ``pcolormesh``/ + ``contour`` receive consistently-shaped synthetic data without needing a + real p2 tensor-basis fixture (see TestGkNodesPsiOverlayRealFixtures).""" + + def _add_psi(self, stub, path): + psi_grid = [np.linspace(0.0, 3.0, 4), np.linspace(-1.0, 1.0, 3)] + psi_values = np.ones((3, 2)) + stub.add(f"{path}sim-psi.gkyl", psi_grid, psi_values) + + def test_pcolormesh_with_colorbar(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + self._add_psi(stub, path) + fig = nodes.nodes("sim", path=path, psi_file="sim-psi.gkyl") + try: + assert len(fig.axes) == 2 + finally: + plt.close(fig) + + def test_contour(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + self._add_psi(stub, path) + fig = nodes.nodes("sim", path=path, psi_file="sim-psi.gkyl", contour=True) + try: + assert len(fig.axes) == 2 + finally: + plt.close(fig) + + def test_single_level_clevels_suppresses_colorbar(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + self._add_psi(stub, path) + fig = nodes.nodes("sim", path=path, psi_file="sim-psi.gkyl", clevels="0.5") + try: + assert len(fig.axes) == 1 + finally: + plt.close(fig) + + def test_absolute_psi_file_override(self, stub, tmp_path): + path = str(tmp_path) + "/" + stub.add(f"{path}sim-nodes.gkyl", [np.arange(3.0)] * 2, _square_nodes()) + abs_psi = f"{path}custom_psi.gkyl" + psi_grid = [np.linspace(0.0, 3.0, 4), np.linspace(-1.0, 1.0, 3)] + stub.add(abs_psi, psi_grid, np.ones((3, 2))) + fig = nodes.nodes("sim", path=path, psi_file=abs_psi) + try: + assert len(fig.axes) == 2 + finally: + plt.close(fig) + + +class TestGkNodesPsiOverlayRealFixtures: + """The psi overlay calls ``GData.interpolate()`` on a real modal DG field + (``nodes`` hardcodes ``poly_order=2``, basis ``"mt"``/tensor, and never + selects a single component before handing the interpolated array straight + to ``pcolormesh``/``contour``) -- skipped loudly since the repo's one + matching-basis fixture, ``tests/test_data/generated/2d_mt_p2.gkyl``, is a + 9-component demo field (no shipped poloidal-flux fixture is single- + component), which ``pcolormesh``/``contour`` cannot render directly.""" + + def test_psi_overlay_needs_single_component_p2_tensor_fixture(self): + import postgkyl as pg + + candidate = os.path.join(GENERATED, "2d_mt_p2.gkyl") + if not os.path.exists(candidate): + pytest.skip(f"no p2 tensor-basis 2-D fixture at '{candidate}'.") + num_comps = pg.load(candidate).num_comps + if num_comps == 1: + pytest.fail("fixture is now single-component -- wire up a real " + "psi-overlay assertion here") + pytest.skip( + f"'{candidate}' has {num_comps} components; nodes(psi_file=...) " + "never selects a single component before pcolormesh/contour, so " + "this fixture cannot exercise that path meaningfully. See " + "TestGkNodesSynthetic for the node-plotting coverage instead.") diff --git a/tests/test_diagnostics_programs_particle_balance.py b/tests/test_diagnostics_programs_particle_balance.py new file mode 100644 index 00000000..b10533ab --- /dev/null +++ b/tests/test_diagnostics_programs_particle_balance.py @@ -0,0 +1,261 @@ +"""Tests for ``postgkyl.diagnostics.gk.particle_balance``. + +See ``test_diagnostics_programs_energy_balance.py`` for the shared testing +strategy (no ``tests_bak`` corpus exists for this app; the repo ships no +multi-file gyrokinetic particle-balance fixture set, so the full figure path +is exercised against synthetic per-file datasets stubbed through +``utils.GData``). + +Run: PYTHONPATH=src pytest tests/test_diagnostics_programs_particle_balance.py -v +""" + +from __future__ import annotations + +import importlib +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from postgkyl.diagnostics.gk import utils as gk_utils + +pb = importlib.import_module("postgkyl.diagnostics.gk.particle_balance") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") + + +class _FakeGData: + + def __init__(self, grid, values, ctx=None): + self._grid = grid + self._values = values + self.ctx = ctx or {} + + def get_grid(self): + return self._grid + + def get_values(self): + return self._values + + +class _StubFiles: + + def __init__(self, tmp_path, monkeypatch): + self._registry: dict[str, _FakeGData] = {} + monkeypatch.setattr(gk_utils, "GData", self._dispatch) + + def _dispatch(self, file_name): + return self._registry[file_name] + + def add(self, file_name: str, time, values) -> None: + open(file_name, "w").close() + self._registry[file_name] = _FakeGData([np.asarray(time)], + np.asarray(values)) + + +@pytest.fixture +def stub(tmp_path, monkeypatch): + return _StubFiles(tmp_path, monkeypatch) + + +def _build_sim(stub, + tmp_path, + name="sim", + species="ion", + *, + with_src=True, + with_bflux=True, + n=5): + path = str(tmp_path) + "/" + # A dynvector's grid is exactly one time stamp per recorded sample, not + # N+1 cell edges like a field file (see io/gkyl_reader.py's _read_t2_v1). + time = np.linspace(0.0, 1.0, n) + + # 2 components (M0, M1): a single-component array would collapse to 1-D + # under np.squeeze, breaking the `v[:, _DENSITY_MOMENT]` indexing every + # integrated-moments file family needs. + fdot_vals = np.zeros((n, 2)) + fdot_vals[:, 0] = np.linspace(1.0, 2.0, n) + stub.add(f"{path}{name}-{species}_fdot_integrated_moms.gkyl", time, fdot_vals) + + if with_src: + src_vals = np.zeros((n, 2)) + src_vals[:, 0] = 0.1 + stub.add(f"{path}{name}-{species}_source_integrated_moms.gkyl", time, + src_vals) + if with_bflux: + bflux_vals = np.zeros((n, 2)) + bflux_vals[:, 0] = 0.05 + stub.add( + f"{path}{name}-{species}_bflux_xlower_integrated_HamiltonianMoments.gkyl", + time, bflux_vals) + return path + + +class TestParticleBalanceErrorPure: + + def test_formula(self): + fdot = np.array([2.0, 3.0]) + src = np.array([1.0, 1.0]) + bflux = np.array([0.5, 0.5]) + err = pb.particle_balance_error(fdot, src, bflux) + np.testing.assert_allclose(err, src - bflux - fdot) + + +class TestAccumulatePure: + + def test_first_use_copies_not_aliases(self): + a = np.array([1.0, 2.0]) + out = pb._accumulate(None, a) + out[0] = 99.0 + assert a[0] == 1.0 + + def test_accumulates_sum(self): + out = pb._accumulate(np.array([1.0, 2.0]), np.array([3.0, 4.0])) + np.testing.assert_allclose(out, [4.0, 6.0]) + + +class TestResolvePure: + + def test_no_override_uses_default(self): + assert pb._resolve("/p/", None, "default.gkyl", 0) == "default.gkyl" + + def test_override_substitutes_block(self): + assert pb._resolve("/p/", "custom_*.gkyl", "unused", + 3) == "/p/custom_3.gkyl" + + +class TestGkParticleBalanceSynthetic: + + def test_full_path_with_src_and_bflux(self, stub, tmp_path): + path = _build_sim(stub, tmp_path) + fig, traces = pb.particle_balance("sim", "ion", path=path) + try: + assert traces.src is not None + assert traces.bflux_tot is not None + assert traces.mom_err is not None + assert traces.time.shape[0] == 5 + finally: + plt.close(fig) + + def test_missing_source_and_bflux(self, stub, tmp_path): + path = _build_sim(stub, tmp_path, with_src=False, with_bflux=False) + fig, traces = pb.particle_balance("sim", "ion", path=path) + try: + assert traces.src is None + assert traces.bflux_tot is None + finally: + plt.close(fig) + + def test_relative_error_branch(self, stub, tmp_path): + path = _build_sim(stub, tmp_path) + n = 5 + time = np.linspace(0.0, 1.0, n) + f_vals = np.zeros((n, 2)) + f_vals[:, 0] = 10.0 + stub.add(f"{path}sim-ion_integrated_moms.gkyl", time, f_vals) + dt_time = np.linspace(0.0, 1.0, n - 1) + dt_vals = np.full((n - 1, 1), 0.2) + stub.add(f"{path}sim-dt.gkyl", dt_time, dt_vals) + + fig, traces = pb.particle_balance("sim", + "ion", + path=path, + relative_error=True) + try: + assert traces.mom_err is None + assert traces.mom_err_norm is not None + assert traces.mom_err_norm.shape[0] == n - 1 + finally: + plt.close(fig) + + def test_relative_error_absy_saveas_and_show(self, stub, tmp_path): + """Covers ``absy`` wrapping the relative-error ylabel in ``||`` together + with ``saveas``/``show`` (default ``ylabel_string`` is non-empty on this + branch, unlike the absolute-error branch's ``""`` default).""" + path = _build_sim(stub, tmp_path) + n = 5 + time = np.linspace(0.0, 1.0, n) + f_vals = np.zeros((n, 2)) + f_vals[:, 0] = 10.0 + stub.add(f"{path}sim-ion_integrated_moms.gkyl", time, f_vals) + dt_time = np.linspace(0.0, 1.0, n - 1) + dt_vals = np.full((n - 1, 1), 0.2) + stub.add(f"{path}sim-dt.gkyl", dt_time, dt_vals) + + out_path = str(tmp_path / "out.png") + fig, traces = pb.particle_balance("sim", + "ion", + path=path, + relative_error=True, + absy=True, + show=True, + saveas=out_path) + try: + assert traces.mom_err_norm is not None + assert os.path.exists(out_path) + finally: + plt.close(fig) + + def test_missing_required_fdot_file_raises(self, stub, tmp_path): + path = str(tmp_path) + "/" + with pytest.raises(FileNotFoundError, match="fdot_integrated_moms"): + pb.particle_balance("sim", "ion", path=path) + + def test_bflux_override_and_absy_logy(self, stub, tmp_path): + path = _build_sim(stub, tmp_path, with_bflux=False) + n = 5 + time = np.linspace(0.0, 1.0, n) + override_vals = np.zeros((n, 2)) + override_vals[:, 0] = -0.05 + override_name = f"{path}custom_bflux.gkyl" + stub.add(override_name, time, override_vals) + + fig, traces = pb.particle_balance( + "sim", + "ion", + path=path, + bflux_files={"xlower": "custom_bflux.gkyl"}, + absy=True, + logy=True) + try: + assert traces.bflux_tot is not None + np.testing.assert_allclose(traces.bflux_tot, -0.05) + finally: + plt.close(fig) + + def test_multiblock_sums_over_blocks(self, stub, tmp_path): + path = str(tmp_path) + "/" + n = 4 + time = np.linspace(0.0, 1.0, n) + for block in (0, 1): + fdot_vals = np.zeros((n, 2)) + fdot_vals[:, 0] = 1.0 + stub.add(f"{path}sim_b{block}-ion_fdot_integrated_moms.gkyl", time, + fdot_vals) + fig, traces = pb.particle_balance("sim", "ion", path=path, multib="0,1") + try: + # Two blocks, each contributing fdot=1.0, sum to 2.0 everywhere. + np.testing.assert_allclose(traces.fdot, 2.0) + finally: + plt.close(fig) + + +class TestGkParticleBalanceRealFixtures: + + def test_real_fixture_particle_balance(self): + required = ("_fdot_integrated_moms.gkyl", ) + if not any( + any(f.endswith(suffix) for f in os.listdir(DATA)) + for suffix in required): + pytest.skip( + "tests/test_data ships no gyrokinetic particle-balance file family " + "(needs e.g. '-_fdot_integrated_moms.gkyl'); see " + "TestGkParticleBalanceSynthetic for full-path coverage against " + "stubbed data instead.") + pytest.fail("fixture files appeared -- wire up a real-data assertion here") diff --git a/tests/test_diagnostics_programs_trajectory.py b/tests/test_diagnostics_programs_trajectory.py new file mode 100644 index 00000000..22125ed8 --- /dev/null +++ b/tests/test_diagnostics_programs_trajectory.py @@ -0,0 +1,189 @@ +"""Tests for ``postgkyl.diagnostics.vm.trajectory``. + +Ported from ``src_bak/postgkyl/apps/trajectory.py`` (no ``tests_bak`` corpus +exists for this app). A Gkeyll dynvector's grid holds exactly one time stamp +per recorded sample (``io/gkyl_reader.py``'s ``_read_t2_v1``: ``grid[0]`` +has the same length as ``values.shape[0]``) -- unlike a *field* file's +``num_cells + 1`` edge convention. ``postgkyl.io.write`` only emits +file_type == 1 (field) ``.gkyl`` files, so a real write -> reload round trip +reads back through a field file with no basis metadata; ``GDataState`` +defaults that case to p0 serendipity/nodal and re-expresses the grid as +cell centers, which lines back up with the dynvector convention (one point +per sample). Most trajectory fixtures here still build the ``GDataState`` +directly (the same technique ``tests/test_io_writer.py``'s ``_make_state`` +uses); one test explicitly exercises the ``io.write`` round trip to confirm +the convention matches after reload. + +Run: PYTHONPATH=src pytest tests/test_diagnostics_programs_trajectory.py -v +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from postgkyl import io +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.diagnostics.vm import trajectory as traj + + +def _make_trajectory(num_pos=10, *, velocity=False, seed=0): + """A synthetic dynvector-shaped trajectory: ``grid[0]`` has exactly + ``num_pos`` time stamps, matching ``values.shape[0]``.""" + rng = np.random.default_rng(seed) + time = np.linspace(0.0, 1.0, num_pos) + ncomp = 6 if velocity else 3 + values = rng.uniform(-1.0, 1.0, size=(num_pos, ncomp)) + d = GDataState() + d.push([time], values) + return d + + +class TestMasked: + + def test_no_bounds_passthrough(self): + coord = np.array([1.0, 2.0, 3.0]) + out = traj._masked(coord, None, None) + np.testing.assert_allclose(out, coord) + + def test_lower_bound_masks_below(self): + coord = np.array([1.0, 2.0, 3.0]) + out = traj._masked(coord, 1.5, None) + assert np.isnan(out[0]) + np.testing.assert_allclose(out[1:], [2.0, 3.0]) + + def test_upper_bound_masks_above(self): + coord = np.array([1.0, 2.0, 3.0]) + out = traj._masked(coord, None, 2.5) + np.testing.assert_allclose(out[:2], [1.0, 2.0]) + assert np.isnan(out[2]) + + def test_both_bounds(self): + coord = np.array([1.0, 2.0, 3.0]) + out = traj._masked(coord, 1.5, 2.5) + assert np.isnan(out[0]) + np.testing.assert_allclose(out[1], 2.0) + assert np.isnan(out[2]) + + +class TestTrajectoryRaises: + + def test_no_datasets_raises(self): + with pytest.raises(ValueError, match="at least one dataset"): + traj.trajectory() + + +class TestTrajectorySynthetic: + + def test_frame_count_matches_samples(self): + d = _make_trajectory(num_pos=8) + anim = traj.trajectory(d) + try: + assert anim._save_count == 8 + finally: + plt.close(anim._fig) + + def test_numframes_subsamples(self): + d = _make_trajectory(num_pos=20) + anim = traj.trajectory(d, numframes=5) + try: + assert anim._save_count == 5 + finally: + plt.close(anim._fig) + + def test_first_frame_renders_without_error(self): + d = _make_trajectory(num_pos=6, velocity=True) + anim = traj.trajectory(d, no_velocity=False) + try: + fig = anim._fig + ax = fig.axes[0] + traj._update(0, ax, (d, ), 1, False, None, None, None, None, None, None) + assert ax.get_title().startswith("T:") + finally: + plt.close(anim._fig) + + def test_last_frame_uses_final_dt_branch(self): + """When ``t_idx + leap`` runs past the end of the trace, the velocity + vector uses ``time[-1] - time[t_idx]`` instead of indexing out of + bounds.""" + d = _make_trajectory(num_pos=4, velocity=True) + fig = plt.figure() + ax = fig.add_subplot(111, projection="3d") + try: + traj._update(3, ax, (d, ), 1, False, None, None, None, None, None, None) + finally: + plt.close(fig) + + def test_multiple_datasets_overlaid(self): + d1 = _make_trajectory(num_pos=6, seed=1) + d2 = _make_trajectory(num_pos=6, seed=2) + anim = traj.trajectory(d1, d2) + try: + assert anim._save_count == 6 + finally: + plt.close(anim._fig) + + def test_axis_bounds_mask_points(self): + d = _make_trajectory(num_pos=6) + anim = traj.trajectory(d, + xmin=-0.5, + xmax=0.5, + ymin=-0.5, + ymax=0.5, + zmin=-0.5, + zmax=0.5) + try: + assert anim._save_count == 6 + finally: + plt.close(anim._fig) + + def test_fixaspect_and_view_angles(self): + d = _make_trajectory(num_pos=5) + anim = traj.trajectory(d, fixaspect=True, elevation=30.0, azimuth=45.0) + try: + assert anim._save_count == 5 + finally: + plt.close(anim._fig) + + +class TestTrajectoryViaIoWriter: + """Exercises the ``io.write`` round trip the instruction file suggests. + + The written file is a field file with no basis metadata, so on reload + ``GDataState`` defaults it to p0 serendipity/nodal and re-expresses the + ``num_cells + 1`` edge grid as cell centers -- which lines back up + one-to-one with the dynvector convention (``len(grid[0]) == values.shape[0]``). + + Single-component only: the compiled reader (``gpython.rio.read_field``, tried + first whenever the shim is available) fails on *any* multi-component + ``.gkyl`` field this writer produces -- + ``PYTHONPATH=src python -c`` reproduction: + ``io.write(state_with_ncomp_2_or_more, ...)`` then re-reading it raises + ``OSError: gpython_read_field failed`` (reproduces even for pre-existing, + layer-agnostic data, e.g. any ``GDataState`` pushed with + ``values.shape[-1] >= 2``; single-component data round-trips fine). That + is a pre-existing limitation in ``gpython``/``io`` (outside this layer's + scope), not something introduced here -- see this layer's report. A real + 3-component trajectory is exercised directly (no disk I/O) by + ``TestTrajectorySynthetic`` instead.""" + + def test_single_component_trajectory_round_trips_and_animates(self, tmp_path): + num_pos = 6 + time_edges = np.linspace(0.0, 1.0, num_pos + 1) + values = np.zeros((num_pos, 1)) + values[:, 0] = np.linspace(0.0, 1.0, num_pos) + d = GDataState() + d.push([time_edges], values) + + out = io.save(d, out_name=str(tmp_path / "traj.gkyl"), extension="gkyl") + + from postgkyl.gdata import GData + with pytest.warns(UserWarning, match="not resolvable"): + reloaded = GData(out) + assert reloaded.grid[0].shape[ + 0] == num_pos + 1 # field convention: N+1 edges + assert reloaded.values.shape[0] == num_pos diff --git a/tests/test_diagnostics_rotations.py b/tests/test_diagnostics_rotations.py new file mode 100644 index 00000000..47bc31f0 --- /dev/null +++ b/tests/test_diagnostics_rotations.py @@ -0,0 +1,120 @@ +"""Tests for postgkyl.diagnostics.mom.rotations -- parrotate/perprotate, folding +the array-math analytic tests (formerly tests_models_rotations.py) with the +verb-level guard/inplace tests (formerly part of tests_ops_physics.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.mom import rotations +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +class TestParrotate: + + def test_u_parallel_to_v_returns_u(self): + u = _make([np.linspace(0.0, 1.0, 3)], + np.array([[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]])) + v = _make([np.linspace(0.0, 1.0, 3)], + np.array([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]])) + out = rotations.parrotate(u, v) + np.testing.assert_allclose(out.values, u.values, atol=1e-12) + + def test_u_perpendicular_to_v_returns_zero(self): + u = _make([np.linspace(0.0, 1.0, 3)], + np.array([[0.0, 1.0, 0.0], [0.0, 2.0, 0.0]])) + v = _make([np.linspace(0.0, 1.0, 3)], + np.array([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]])) + out = rotations.parrotate(u, v) + np.testing.assert_allclose(out.values, np.zeros_like(u.values), atol=1e-12) + + def test_u_oblique_to_v(self): + u = _make([np.array([0.0, 1.0])], np.array([[3.0, 4.0, 0.0]])) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + out = rotations.parrotate(u, v) + np.testing.assert_allclose(out.values[0], [3.0, 0.0, 0.0], atol=1e-12) + + def test_custom_rotate_coords(self): + u = _make([np.array([0.0, 1.0])], np.array([[3.0, 4.0, 0.0]])) + field = _make([np.array([0.0, 1.0])], + np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]])) + out = rotations.parrotate(u, field, coords="3:6") + np.testing.assert_allclose(out.values[0], [3.0, 0.0, 0.0], atol=1e-12) + + def test_grid_passed_through(self): + grid = [np.linspace(0.0, 1.0, 2)] + u = _make(grid, np.array([[1.0, 0.0, 0.0]])) + v = _make(grid, np.array([[1.0, 0.0, 0.0]])) + out = rotations.parrotate(u, v) + np.testing.assert_allclose(out.grid[0], grid[0]) + + def test_mismatched_components_raises(self): + u = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0]])) # only 2 comps + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + with pytest.raises(ValueError, match="three-component"): + rotations.parrotate(u, v) + + def test_inplace_mutates_array(self): + u = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + out = rotations.parrotate(u, v, inplace=True) + assert out is u + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + rotations.parrotate(d, v) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + rotations.parrotate(v, d) + + +class TestPerprotate: + + def test_u_parallel_to_v_gives_zero(self): + u = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + out = rotations.perprotate(u, v) + np.testing.assert_allclose(out.values, np.zeros_like(u.values), atol=1e-12) + + def test_u_perpendicular_to_v_gives_u(self): + u = _make([np.array([0.0, 1.0])], np.array([[0.0, 1.0, 0.0]])) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + out = rotations.perprotate(u, v) + np.testing.assert_allclose(out.values, u.values, atol=1e-12) + + def test_perp_plus_par_equals_u(self): + u = _make([np.array([0.0, 1.0])], np.array([[3.0, 4.0, 0.0]])) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + par = rotations.parrotate(u, v) + perp = rotations.perprotate(u, v) + np.testing.assert_allclose(par.values + perp.values, u.values, atol=1e-12) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + v = _make([np.array([0.0, 1.0])], np.array([[1.0, 0.0, 0.0]])) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + rotations.perprotate(d, v) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + rotations.perprotate(v, d) diff --git a/tests/test_diagnostics_ten_moment.py b/tests/test_diagnostics_ten_moment.py new file mode 100644 index 00000000..d6b27928 --- /dev/null +++ b/tests/test_diagnostics_ten_moment.py @@ -0,0 +1,359 @@ +"""Tests for postgkyl.diagnostics.mom.ten_moment -- 10-moment pressure tensor, +field-aligned pressure diagnostics (p_par, p_perp, agyrotropy), folding the +array-math analytic tests (formerly tests_models_ten_moment.py) with the +verb-level guard/inplace/VARIABLES tests (formerly part of +tests_ops_moments.py / tests_ops_physics.py).""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.diagnostics.mom import ten_moment as tm +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +_G1D = [np.array([0.0, 1.0])] + +_RHO = 1.0 +_VX, _VY, _VZ = 0.5, 0.25, 0.1 +_P_T = 0.4 +_MOM10 = np.array([[ + _RHO, _RHO * _VX, _RHO * _VY, _RHO * _VZ, _P_T + _RHO * _VX**2, + _RHO * _VX * _VY, _RHO * _VX * _VZ, _P_T + _RHO * _VY**2, _RHO * _VY * _VZ, + _P_T + _RHO * _VZ**2 +]]) + + +def _diagonal_pressure(pxx, pyy, pzz): + return _make(_G1D, np.array([[pxx, 0.0, 0.0, pyy, 0.0, pzz]])) + + +def _b(bx, by, bz): + return _make(_G1D, np.array([[bx, by, bz]])) + + +class TestPressureTensorComponents: + + def test_pxx(self): + d = _make(_G1D, _MOM10) + out = tm.pxx(d) + np.testing.assert_allclose(out.values[0, 0], _P_T, rtol=1e-10) + + def test_pxy_pxz_pyz_zero_for_diagonal_flow(self): + d = _make(_G1D, _MOM10) + np.testing.assert_allclose(tm.pxy(d).values[0, 0], 0.0, atol=1e-14) + np.testing.assert_allclose(tm.pxz(d).values[0, 0], 0.0, atol=1e-14) + np.testing.assert_allclose(tm.pyz(d).values[0, 0], 0.0, atol=1e-14) + + def test_pyy(self): + d = _make(_G1D, _MOM10) + out = tm.pyy(d) + np.testing.assert_allclose(out.values[0, 0], _P_T, rtol=1e-10) + + def test_pzz(self): + d = _make(_G1D, _MOM10) + out = tm.pzz(d) + np.testing.assert_allclose(out.values[0, 0], _P_T, rtol=1e-10) + + def test_pressure_tensor_shape_and_diagonal(self): + d = _make(_G1D, _MOM10) + out = tm.pressure_tensor(d) + assert out.values.shape[-1] == 6 + np.testing.assert_allclose(out.values[0, 0], _P_T, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 3], _P_T, rtol=1e-10) + np.testing.assert_allclose(out.values[0, 5], _P_T, rtol=1e-10) + np.testing.assert_allclose(out.values[0, [1, 2, 4]], 0.0, atol=1e-14) + + @needs_gkeyll + def test_pxx_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + tm.pxx(d) + + +class TestPPar: + + def test_b_along_x_pxx_is_p_par(self): + p = _diagonal_pressure(1.0, 0.5, 0.5) + b = _b(1.0, 0.0, 0.0) + out = tm.p_par(p, b) + np.testing.assert_allclose(out.values.flat[0], 1.0, rtol=1e-12) + + def test_b_along_y_pyy_is_p_par(self): + p = _diagonal_pressure(0.5, 2.0, 0.5) + b = _b(0.0, 1.0, 0.0) + out = tm.p_par(p, b) + np.testing.assert_allclose(out.values.flat[0], 2.0, rtol=1e-12) + + def test_b_along_z_pzz_is_p_par(self): + p = _diagonal_pressure(0.5, 0.5, 3.0) + b = _b(0.0, 0.0, 1.0) + out = tm.p_par(p, b) + np.testing.assert_allclose(out.values.flat[0], 3.0, rtol=1e-12) + + def test_isotropic_pressure_p_par_equals_p(self): + p = _diagonal_pressure(2.0, 2.0, 2.0) + b = _b(1.0, 1.0, 0.0) + out = tm.p_par(p, b) + np.testing.assert_allclose(out.values.flat[0], 2.0, rtol=1e-10) + + def test_b_diagonal_gives_average(self): + p = _diagonal_pressure(1.0, 2.0, 0.0) + b = _b(1.0 / np.sqrt(2), 1.0 / np.sqrt(2), 0.0) + out = tm.p_par(p, b) + np.testing.assert_allclose(out.values.flat[0], 1.5, rtol=1e-12) + + def test_inplace_mutates_ptensor(self): + p = _diagonal_pressure(1.0, 0.5, 0.5) + b = _b(1.0, 0.0, 0.0) + out = tm.p_par(p, b, inplace=True) + assert out is p + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + b = _b(1.0, 0.0, 0.0) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + tm.p_par(d, b) + + +class TestPPerp: + + def test_b_along_x_perp_is_average_of_pyy_pzz(self): + p = _diagonal_pressure(1.0, 0.6, 0.4) + b = _b(1.0, 0.0, 0.0) + out = tm.p_perp(p, b) + np.testing.assert_allclose(out.values.flat[0], 0.5, rtol=1e-12) + + def test_isotropic_pressure_perp_equals_par(self): + p = _diagonal_pressure(1.5, 1.5, 1.5) + b = _b(1.0, 0.0, 0.0) + par_out = tm.p_par(p, b) + perp_out = tm.p_perp(p, b) + np.testing.assert_allclose(perp_out.values.flat[0], + par_out.values.flat[0], + rtol=1e-10) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + b = _b(1.0, 0.0, 0.0) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + tm.p_perp(d, b) + + +class TestAgyro: + + @pytest.mark.parametrize("measure", ["frobenius", "swisdak"]) + def test_isotropic_tensor_is_gyrotropic(self, measure): + p = _diagonal_pressure(2.0, 2.0, 2.0) + b = _b(0.0, 0.0, 1.0) + out = tm.agyro(p, b, measure=measure) + np.testing.assert_allclose(out.values, 0.0, atol=1e-10) + + def test_swisdak_case_insensitive(self): + p = _make(_G1D, np.array([[2.0, 0.5, 0.0, 1.0, 0.0, 1.0]])) + b = _b(1.0, 0.0, 0.0) + out1 = tm.agyro(p, b, measure="swisdak") + out2 = tm.agyro(p, b, measure="Swisdak") + np.testing.assert_allclose(out1.values, out2.values) + + def test_frobenius_case_insensitive(self): + p = _make(_G1D, np.array([[2.0, 0.5, 0.0, 1.0, 0.0, 1.0]])) + b = _b(1.0, 0.0, 0.0) + out1 = tm.agyro(p, b, measure="frobenius") + out2 = tm.agyro(p, b, measure="Frobenius") + np.testing.assert_allclose(out1.values, out2.values) + + def test_invalid_measure_raises(self): + p = _diagonal_pressure(1.0, 1.0, 1.0) + b = _b(1.0, 0.0, 0.0) + with pytest.raises(ValueError, match="swisdak.*frobenius"): + tm.agyro(p, b, measure="invalid") + + def test_agyrotropic_swisdak_nonzero(self): + p = _make(_G1D, np.array([[2.0, 0.5, 0.0, 1.0, 0.0, 1.0]])) + b = _b(1.0, 0.0, 0.0) + out = tm.agyro(p, b, measure="swisdak") + assert out.values.flat[0] > 0.0 + + def test_agyrotropic_frobenius_nonzero(self): + p = _make(_G1D, np.array([[2.0, 0.5, 0.0, 1.0, 0.0, 1.0]])) + b = _b(1.0, 0.0, 0.0) + out = tm.agyro(p, b, measure="frobenius") + assert out.values.flat[0] > 0.0 + + def test_default_measure_is_frobenius(self): + p = _make(_G1D, np.array([[2.0, 0.5, 0.0, 1.0, 0.0, 1.0]])) + b = _b(1.0, 0.0, 0.0) + default_out = tm.agyro(p, b) + explicit_out = tm.agyro(p, b, measure="frobenius") + np.testing.assert_allclose(default_out.values, explicit_out.values) + + def test_inplace_mutates_ptensor(self): + p = _diagonal_pressure(2.0, 2.0, 2.0) + b = _b(0.0, 0.0, 1.0) + out = tm.agyro(p, b, inplace=True) + assert out is p + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + b = _b(0.0, 0.0, 1.0) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + tm.agyro(d, b) + + +class TestMomAgyro: + + def _species_and_field(self): + species = _make( + _G1D, np.array([[1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 2.0, 0.0, 2.0]])) + field = _make(_G1D, np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 1.0]])) + return species, field + + def test_isotropic_species_is_gyrotropic(self): + species, field = self._species_and_field() + out = tm.mom_agyro(species, field) + np.testing.assert_allclose(out.values, 0.0, atol=1e-12) + + def test_matches_private_helper(self): + species, field = self._species_and_field() + out = tm.mom_agyro(species, field, measure="swisdak") + _, expected = tm._get_gkyl_10m_agyro(species.grid, + species.values, + field.grid, + field.values, + measure="swisdak") + np.testing.assert_allclose(out.values, expected) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + field = _make(_G1D, np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 1.0]])) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + tm.mom_agyro(d, field) + + +class TestGkyl10mPrivateWrappers: + """The ``_get_gkyl_10m_p_par``/``_get_gkyl_10m_p_perp`` helpers have no + public GData wrapper (the target layout table for this module lists no + 'mom_p_par'/'mom_p_perp' verb, unlike ``mom_agyro``) -- ported directly + against the private array-level functions, matching the old + ``models``-level tests exactly.""" + + @staticmethod + def _species_and_field(): + rho, vx = 1.0, 0.5 + Pxx = 2.0 + rho * vx**2 + Pxy = 0.3 + mom10 = np.array([[rho, rho * vx, 0.0, 0.0, Pxx, Pxy, 0.0, 1.0, 0.0, 1.0]]) + field_vals = np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]]) + g = [np.array([0.0, 1.0])] + return g, mom10, g, field_vals + + def test_p_par_wrapper(self): + sg, sv, fg, fv = self._species_and_field() + _, p_par = tm._get_gkyl_10m_p_par(sg, sv, fg, fv) + np.testing.assert_allclose(p_par.flat[0], 2.0, rtol=1e-10) + + def test_p_perp_wrapper(self): + sg, sv, fg, fv = self._species_and_field() + _, p_perp = tm._get_gkyl_10m_p_perp(sg, sv, fg, fv) + np.testing.assert_allclose(p_perp.flat[0], 1.0, rtol=1e-10) + + +class TestFiveMomentSetFixedAtTenMoments: + + def _tenmoment_state(self): + vals = np.array([[1.0, 2.0, 0.0, 0.0, 6.0, 0.0, 0.0, 3.0, 0.0, 3.0]]) + return _make([np.array([0.0, 1.0])], vals) + + def test_density_reused_from_five_moment(self): + from postgkyl.diagnostics.mom import five_moment as fm + assert tm.density is fm.density + assert tm.xvel is fm.xvel + assert tm.vel is fm.vel + + def test_pressure_uses_num_moms_10(self): + d = self._tenmoment_state() + out = tm.pressure(d) + from postgkyl.diagnostics.mom.five_moment import _get_p + _, expected = _get_p(d.grid, d.values, gas_gamma=5.0 / 3, num_moms=10) + np.testing.assert_allclose(out.values, expected) + + def test_ke_uses_num_moms_10(self): + d = self._tenmoment_state() + out = tm.ke(d) + from postgkyl.diagnostics.mom.five_moment import _get_ke + _, expected = _get_ke(d.grid, d.values, gas_gamma=5.0 / 3, num_moms=10) + np.testing.assert_allclose(out.values, expected) + + def test_temp_uses_num_moms_10(self): + d = self._tenmoment_state() + out = tm.temp(d) + from postgkyl.diagnostics.mom.five_moment import _get_temp + _, expected = _get_temp(d.grid, d.values, gas_gamma=5.0 / 3, num_moms=10) + np.testing.assert_allclose(out.values, expected) + + def test_sound_uses_num_moms_10(self): + d = self._tenmoment_state() + out = tm.sound(d) + from postgkyl.diagnostics.mom.five_moment import _get_sound + _, expected = _get_sound(d.grid, d.values, gas_gamma=5.0 / 3, num_moms=10) + np.testing.assert_allclose(out.values, expected) + + def test_mach_uses_num_moms_10(self): + d = self._tenmoment_state() + out = tm.mach(d) + from postgkyl.diagnostics.mom.five_moment import _get_mach + _, expected = _get_mach(d.grid, d.values, gas_gamma=5.0 / 3, num_moms=10) + np.testing.assert_allclose(out.values, expected) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + tm.density(d) + + @needs_gkeyll + @pytest.mark.parametrize("fn_name", ["ke", "temp", "sound", "mach"]) + def test_all_scalar_quantities_reject_modal_data(self, fn_name): + d = pg.load(F1) + fn = getattr(tm, fn_name) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + fn(d) + + +class TestVariables: + + def test_variables_table_has_exactly_the_old_tenmoment_vocabulary(self): + assert set(tm.VARIABLES) == { + "density", "xvel", "yvel", "zvel", "vel", "pressure", "ke", "temp", + "sound", "mach", "pressureTensor", "pxx", "pxy", "pxz", "pyy", "pyz", + "pzz" + } + + def test_variables_table_maps_to_public_functions(self): + assert tm.VARIABLES["pressureTensor"] is tm.pressure_tensor + assert tm.VARIABLES["pxx"] is tm.pxx + assert tm.VARIABLES["density"] is tm.density diff --git a/tests/test_docs_build.py b/tests/test_docs_build.py new file mode 100644 index 00000000..85ab7d02 --- /dev/null +++ b/tests/test_docs_build.py @@ -0,0 +1,173 @@ +"""Exercise the publishable docs, including the downloadable example bundle.""" + +from __future__ import annotations + +import json +import os +import pickle +from pathlib import Path +import runpy +import subprocess +import sys +import zipfile + +import pytest + +from postgkyl import gpython +from postgkyl.cli.app import cli + +pytest.importorskip("sphinx", reason="install the docs extra for website tests") +pytestmark = pytest.mark.skipif(not gpython.available(), + reason="needs compiled Gkeyll") +ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture(scope="module") +def documentation(tmp_path_factory): + destination = tmp_path_factory.mktemp("docs") + source = destination / "source" + subprocess.run([ + sys.executable, + str(ROOT / "scripts/build_docs.py"), "--output", + str(source) + ], + cwd=ROOT, + check=True, + capture_output=True, + text=True) + subprocess.run([ + sys.executable, "-m", "sphinx", "-W", "--keep-going", "-b", "html", "-c", + str(ROOT / "docs"), + str(source), + str(destination / "html") + ], + check=True, + capture_output=True, + text=True) + return destination + + +def test_published_inventory_and_navigation(documentation): + inventory = json.loads((documentation / "source/build-info.json").read_text()) + assert set(inventory["commands"]) == set(cli.commands) + index = (documentation / "html/index.html").read_text() + assert 'href="examples.html"' in index + with (documentation / + "html/.doctrees/environment.pickle").open("rb") as stream: + environment = pickle.load(stream) + assert environment.toctree_includes["index"] == [ + "installation", "examples", "reference/api", "reference/cli", "concepts", + "reference/quantities", "contributing", "provenance" + ] + examples = environment.toctree_includes["examples"] + assert {"cli-tutorial", "interface-equivalence"} <= set(examples) + pairs = [] + for page in examples: + assert (documentation / f"html/{page}.html").is_file() + source = (documentation / f"source/{page}.rst").read_text() + includes = [ + line for line in source.splitlines() + if line.startswith(".. include:: _pairs/") + ] + assert len(includes) <= 1, page + pairs.extend(line.split("_pairs/")[1] for line in includes) + assert sorted(pairs) == sorted( + path.name for path in (documentation / "source/_pairs").glob("*.inc")) + assert 'href="reference/cli.html"' in index + api = (documentation / "html/reference/api.html").read_text() + assert "postgkyl.load" in api + for path, page in inventory["api_pages"].items(): + assert (documentation / f"html/reference/{page}.html").is_file(), path + load_page = inventory["api_pages"]["postgkyl.load"] + assert "poly_order" in (documentation / + f"html/reference/{load_page}.html").read_text() + pressure_page = inventory["api_pages"][ + "postgkyl.diagnostics.mom.five_moment.pressure"] + assert "gas_gamma" in (documentation / + f"html/reference/{pressure_page}.html").read_text() + comparison = json.loads( + (documentation / "source/figures/interface-comparison.json").read_text()) + commands = json.loads((ROOT / "examples/figure_commands.json").read_text()) + assert set(comparison) == { + name + for outputs in commands.values() + for name in outputs + } + for prefix in ("python", "cli"): + assert (documentation / + f"html/interactive/{prefix}-08_surface.html").is_file() + assert (documentation / + f"html/interactive/{prefix}-08_volume.html").is_file() + + +def test_downloaded_examples_run_outside_repository(documentation, tmp_path): + with zipfile.ZipFile(documentation / + "source/downloads/postgkyl-examples.zip") as bundle: + bundle.extractall(tmp_path) + env = { + **os.environ, "MPLBACKEND": "Agg", + "PGKYL_EXAMPLE_OUTPUT": str(tmp_path / "figures") + } + subprocess.run([ + sys.executable, + str(tmp_path / "examples/compare_interfaces.py"), "--output", + str(tmp_path / "figures") + ], + cwd=tmp_path, + env=env, + check=True, + capture_output=True, + text=True) + assert (tmp_path / "figures/06_growth.png").stat().st_size > 0 + assert (tmp_path / "figures/05_gk_rz.png").stat().st_size > 0 + + +def test_comparison_rejects_changed_pixels(tmp_path): + from PIL import Image + first, second = tmp_path / "python.png", tmp_path / "cli.png" + Image.new("RGB", (4, 4), "white").save(first) + changed = Image.new("RGB", (4, 4), "white") + changed.putpixel((1, 1), (0, 0, 0)) + changed.save(second) + compare = runpy.run_path(str( + ROOT / "examples/compare_interfaces.py"))["compare_outputs"] + with pytest.raises(AssertionError): + compare(first, second) + + +def test_plotly_comparison_ignores_only_html_id(tmp_path): + first, second = tmp_path / "python.html", tmp_path / "cli.html" + first.write_text( + 'Plotly.newPlot("aaaa", [{"z":[1,2]}], {"title":"Density"}, {})') + second.write_text( + 'Plotly.newPlot("bbbb", [{"z":[1,2]}], {"title":"Density"}, {})') + compare = runpy.run_path(str( + ROOT / "examples/compare_interfaces.py"))["compare_outputs"] + assert "Identical Plotly" in compare(first, second) + second.write_text(second.read_text().replace('[1,2]', '[1,3]')) + with pytest.raises(AssertionError, match="Plotly"): + compare(first, second) + + +def test_animation_comparison_checks_timing(tmp_path): + from PIL import Image + first, second = tmp_path / "python.gif", tmp_path / "cli.gif" + frames = [Image.new("RGB", (4, 4), color) for color in ("white", "black")] + for path, duration in ((first, 100), (second, 200)): + frames[0].save(path, + save_all=True, + append_images=frames[1:], + duration=duration) + compare = runpy.run_path(str( + ROOT / "examples/compare_interfaces.py"))["compare_outputs"] + with pytest.raises(AssertionError, match="timings"): + compare(first, second) + + +def test_preparation_preserves_unowned_directory(tmp_path): + existing = tmp_path / "notes.txt" + existing.write_text("Keep this file") + prepare = runpy.run_path(str(ROOT / "scripts/build_docs.py"))["prepare"] + with pytest.raises(ValueError, match="unowned"): + prepare(ROOT, tmp_path) + assert existing.read_text() == "Keep this file" diff --git a/tests/test_documentation.py b/tests/test_documentation.py new file mode 100644 index 00000000..021c459e --- /dev/null +++ b/tests/test_documentation.py @@ -0,0 +1,132 @@ +"""Contracts keeping Python hover help and generated CLI help in sync.""" + +from __future__ import annotations + +import ast +import inspect +from pathlib import Path +import pydoc + +import postgkyl as pg +from postgkyl.cli.app import COMMANDS, MODELS +from postgkyl.cli.docstrings import parse_docstring +from postgkyl.gdata.gdata import GData + + +def _public_documented_members(cls) -> dict[str, object]: + members = {} + for name in dir(cls): + if name.startswith("_"): + continue + value = getattr(cls, name) + if callable(value) or isinstance(value, property): + members[name] = value + return members + + +def _function_source_doc(function) -> str | None: + path = Path(inspect.getsourcefile(function)) + tree = ast.parse(path.read_text(), path) + matches = [ + node for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == function.__name__ + ] + assert len( + matches) == 1, f"could not locate {function.__qualname__} in {path}" + return ast.get_docstring(matches[0], clean=True) + + +def test_public_python_surfaces_have_hover_documentation(): + missing = [] + for name in pg.__all__: + value = getattr(pg, name) + if callable(value) and not inspect.getdoc(value): + missing.append(f"postgkyl.{name}") + if inspect.isfunction(value): + assert _function_source_doc(value) == inspect.getdoc(value) + for cls in (pg.GData, pg.GDataGroup): + for name, value in _public_documented_members(cls).items(): + if not inspect.getdoc(value): + missing.append(f"{cls.__name__}.{name}") + assert missing == [] + + +def test_fluent_operations_are_static_aliases_to_documented_functions(): + aliases = { + name: value + for name, value in GData.__dict__.items() + if inspect.isfunction(value) and (value.__module__.startswith( + "postgkyl.operations") or value.__module__ == "postgkyl.io.writer") + } + assert aliases + + path = Path(inspect.getsourcefile(GData)) + tree = ast.parse(path.read_text(), path) + class_node = next(node for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "GData") + class_assignments = { + target.id + for node in class_node.body if isinstance(node, ast.Assign) + for target in node.targets if isinstance(target, ast.Name) + } + assert set(aliases) <= class_assignments + + instance = GData() + for name, function in aliases.items(): + method = getattr(instance, name) + assert GData.__dict__[name] is function + assert inspect.getdoc(method) == inspect.getdoc(function) + assert tuple(inspect.signature(method).parameters) == tuple( + inspect.signature(function).parameters)[1:] + + +def test_shared_functional_and_fluent_spellings_are_one_object(): + exceptions = {"load", "plot", "val2coord"} + shared = set(pg.__all__) & set(GData.__dict__) - exceptions + paired = { + name + for name in shared if inspect.isfunction(getattr(pg, name)) + and inspect.isfunction(GData.__dict__[name]) + } + assert paired + for name in paired: + assert GData.__dict__[name] is getattr(pg, name) + + +def test_cli_help_is_lowered_from_source_docstrings(): + commands = {model.name: command for model, command in zip(MODELS, COMMANDS)} + for model in MODELS: + parsed = parse_docstring( + model.canonical, + required=set(inspect.signature(model.canonical).parameters), + signature_names=set(inspect.signature(model.canonical).parameters)) + source_doc = _function_source_doc(model.canonical) + + assert source_doc == inspect.getdoc(model.canonical) + assert "Value for ``" not in source_doc + assert model.help == parsed.summary + assert model.long_help == parsed.long_help + + command = commands[model.name] + assert command.short_help == parsed.summary + assert command.help == parsed.long_help + click_parameters = { + parameter.name: parameter + for parameter in command.params + } + for parameter in model.parameters: + if not parameter.injected: + assert click_parameters[parameter.name].help == parsed.parameters[ + parameter.name] + + +def test_python_help_renders_for_function_and_bound_method(): + summary = "Interpolate DG (modal/nodal) data onto a uniform evaluation mesh." + assert summary in pydoc.render_doc(pg.interpolate) + assert summary in pydoc.render_doc(GData().interpolate) + + +def test_distribution_marks_inline_types_for_editor_tools(): + marker = Path(pg.__file__).with_name("py.typed") + assert marker.is_file() diff --git a/tests/test_examples.py b/tests/test_examples.py new file mode 100644 index 00000000..fb600e01 --- /dev/null +++ b/tests/test_examples.py @@ -0,0 +1,102 @@ +"""The examples in ``examples/`` are the user-facing tutorial (README + +narrated scripts + a CLI walkthrough). This file is what keeps that tutorial +honest: every script's own ``assert`` statements run for real (via +``runpy``, so a broken example fails here with the same traceback a user +would see), and every ``pgkyl ...`` line quoted in ``examples/cli_tutorial.md`` +is replayed through the real CLI. Nothing here re-describes the examples -- +it just executes the one copy of them that already exists. +""" + +from __future__ import annotations + +import glob +import os +import re +import runpy +import shlex + +import matplotlib +import pytest +from click.testing import CliRunner + +matplotlib.use("Agg") + +from postgkyl import gpython +from postgkyl.cli.app import cli + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +EXAMPLES = os.path.join(ROOT, "examples") +SCRIPTS = sorted( + script for script in glob.glob(os.path.join(EXAMPLES, "scripts", "*.py")) + if not os.path.basename(script).startswith("_")) +TUTORIAL = os.path.join(EXAMPLES, "cli_tutorial.md") + + +def _extract_cli_commands(markdown_path): + """Pull every ``pgkyl ...`` invocation out of the ``` ```bash``` ``` fences + in a tutorial markdown file, joining ``\\``-continued lines.""" + with open(markdown_path) as fh: + text = fh.read() + + commands = [] + for block in re.findall(r"```bash\n(.*?)```", text, re.DOTALL): + pending = [] + for line in block.strip("\n").splitlines(): + line = line.rstrip() + if not line: + continue + if line.endswith("\\"): + pending.append(line[:-1]) + continue + pending.append(line) + full = " ".join(pending).strip() + pending = [] + if full.startswith("pgkyl "): + commands.append(full[len("pgkyl "):]) + return commands + + +CLI_COMMANDS = _extract_cli_commands(TUTORIAL) + + +class TestTutorialScripts: + """Every script under ``examples/scripts/`` runs to completion and its + internal assertions pass -- exercised in-process via ``runpy`` so an + ``AssertionError`` inside the example surfaces as this test's failure.""" + + @needs_gkeyll + @pytest.mark.parametrize("script", + SCRIPTS, + ids=[os.path.basename(s) for s in SCRIPTS]) + def test_script_runs_clean(self, script, tmp_path, monkeypatch): + monkeypatch.setenv("PGKYL_EXAMPLE_OUTPUT", str(tmp_path)) + monkeypatch.syspath_prepend(os.path.dirname(script)) + runpy.run_path(script, run_name="__main__") + + +class TestCliTutorial: + """Every ``pgkyl ...`` line quoted in ``examples/cli_tutorial.md`` actually + runs, from the repository root (the fixture paths in the tutorial are + written relative to it, exactly as a reader would type them).""" + + def test_tutorial_has_commands(self): + # A parsing regression (e.g. a fence typo) would otherwise silently + # leave the parametrized test below with zero cases -- a green suite + # that covers nothing. + assert len(CLI_COMMANDS) >= 8 + + @needs_gkeyll + @pytest.mark.parametrize("command", CLI_COMMANDS) + def test_command_succeeds(self, command, tmp_path, monkeypatch): + # Symlink 'tests/' into an isolated cwd: the tutorial's relative input + # paths (``tests/test_data/...``) resolve, but any file the command + # writes (out.png, distf.npy, ...) lands in tmp_path, not the repo. + os.symlink(os.path.join(ROOT, "tests"), os.path.join(tmp_path, "tests")) + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, shlex.split(command)) + assert result.exit_code == 0, ( + f"`pgkyl {command}` failed:\n{result.output}\n{result.exception}") diff --git a/tests/test_gdata_fluent.py b/tests/test_gdata_fluent.py new file mode 100644 index 00000000..75046d28 --- /dev/null +++ b/tests/test_gdata_fluent.py @@ -0,0 +1,483 @@ +"""Tests for the fluent surface (layer 11 -- api): every ``operations`` verb from +layers 07-09 as a ``GData`` method (or, for the multi-dataset verbs with no +single ``self``, a module-level function in ``api.verbs``), the fluent +``api.group.GDataGroup`` that broadcasts verbs over its members, and the +facade re-exports. + +Physics diagnostics are deliberately not fluent methods. Domain-specific data +transformations such as ``gk_rz`` are operations and do belong on the fluent +surface alongside domain-independent core verbs. +""" + +from __future__ import annotations + +import base64 +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations, render +from postgkyl.gdata.gdatagroup import GDataGroup as ApiGDataGroup +from postgkyl.gdata import verbs as api_verbs +from postgkyl.gdatastate.gdatastategroup import GDataStateGroup as CoreGDataStateGroup +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GEN = os.path.join(DATA, "generated") +F1D = os.path.join(GEN, "1d_ms_p1.gkyl") +F2D_VEC = os.path.join(GEN, + "2d_c2p_rot45_ms_p1.gkyl") # 2 comps after interpolate + + +@pytest.fixture(autouse=True) +def _close_figs(): + plt.close("all") + yield + plt.close("all") + + +class MyData(pg.GData): + """A ``GData`` subclass, used to verify subclass propagation through every + fluent method (the ``_result``/``type(self)`` contract).""" + + +def _make(cls, grid, values, **ctx): + d = cls(ctx=ctx or None) + d.push(list(grid), values) + return d + + +def _line(cls=MyData, tag: str = "default", value: float = 1.0, n: int = 5): + grid = [np.linspace(0.0, 1.0, n + 1)] + return _make(cls, grid, np.full((n, 1), value), tag=tag) + + +# ============================================================ method roster +# The full public fluent inventory, split by which object owns each operation. +# GDataGroup broadcasts every single-dataset verb and explicitly implements the +# operations that act on the group as a whole. Every multi-dataset operation +# also has a functional spelling on the top-level ``pg`` facade. +INSTANCE_VERBS = [ + "load", "interpolate", "local_poly", "gk_rz", "select", "plot", "plotly", + "pyvista", "save", "mul", "div", "integrate", "average", + "eval_at_coord_proj", "to_modal", "to_nodal", "to_quad", "apply", "fft", + "magsq", "mask", "val2coord", "extract_input", "fit", "differentiate", "map" +] +GROUP_VERBS = ["sort", "collect", "evaluate", "animate", "plotly_animate"] +MODULE_VERBS = GROUP_VERBS + ["relchange"] + + +class TestMethodInventory: + + def test_every_instance_verb_exists_and_is_callable(self): + data = _line() + for name in INSTANCE_VERBS: + assert hasattr(pg.GData, name), f"GData has no {name!r} method" + assert callable(getattr(pg.GData, name)) + assert callable(getattr(data, + name)), f"GData object has no {name!r} method" + + def test_every_instance_verb_is_available_on_a_group(self): + group = ApiGDataGroup([_line()]) + for name in INSTANCE_VERBS: + assert callable(getattr( + group, name)), (f"GDataGroup object has no {name!r} method") + + def test_every_group_verb_is_a_method_and_a_pg_function(self): + group = ApiGDataGroup([_line()]) + for name in GROUP_VERBS: + assert callable(getattr( + group, name)), (f"GDataGroup object has no {name!r} method") + assert callable(getattr(pg, name)), f"postgkyl has no {name!r} function" + + def test_every_module_verb_exists_in_api_verbs_and_pg(self): + for name in MODULE_VERBS: + assert hasattr(api_verbs, name), f"api.verbs has no {name!r} function" + assert callable(getattr(api_verbs, name)) + assert getattr(pg, name) is getattr(api_verbs, name) + + def test_grid_is_deliberately_not_a_fluent_method(self): + """``operations.grid`` has no fluent spelling: ``GData.grid`` must stay the + inherited axis-edge-array *property* (see api/gdata.py's note), not a + verb method -- otherwise every other verb reading ``data.grid`` would + silently break.""" + d = _line() + assert isinstance(d.grid, list) + assert not callable(d.grid) + assert hasattr(operations, "grid") and callable(operations.grid) + + +# ==================================================== subclass propagation +class TestSubclassPropagation: + + def test_fft(self): + d = _line(value=1.0, n=16) + out = d.fft() + assert isinstance(out, MyData) + out_psd = d.fft(psd=True) + assert isinstance(out_psd, MyData) + + def test_magsq(self): + d = _make(MyData, [np.linspace(0.0, 1.0, 5)], + np.tile([1.0, 2.0, 3.0], (4, 1))) + out = d.magsq() + assert isinstance(out, MyData) + + def test_mask(self): + d = _make(MyData, [np.linspace(0.0, 1.0, 6)], np.arange(5.0)[:, np.newaxis]) + out = d.mask(lower=2.0) + assert isinstance(out, MyData) + + def test_relchange(self): + grid = [np.linspace(0.0, 1.0, 5)] + ref = _make(MyData, grid, np.full((4, 1), 2.0)) + cur = _make(MyData, grid, np.full((4, 1), 3.0)) + out = api_verbs.relchange(ref, cur) + assert isinstance(out, MyData) + + def test_val2coord_returns_fluent_group_of_the_subclass(self): + d = _make(MyData, [np.arange(5.0)], np.arange(15.0).reshape(5, 3)) + group = d.val2coord(x="0", y="1,2") + assert isinstance(group, ApiGDataGroup) + assert len(group) == 2 + for member in group: + assert isinstance(member, MyData) + + def test_extract_input_returns_a_plain_string(self): + d = _line() + assert d.extract_input() == "" + text = "title = my sim\n" + encoded = base64.encodebytes(text.encode("utf-8")).decode("utf-8") + d2 = _make(MyData, [np.linspace(0.0, 1.0, 3)], + np.ones((2, 1)), + input_file=encoded) + assert d2.extract_input() == text + + def test_fit(self): + edges = np.linspace(0.0, 1.0, 21) + centers = 0.5 * (edges[:-1] + edges[1:]) + y = 2.0 * centers + 1.0 + d = _make(MyData, [edges], y[:, np.newaxis]) + out = d.fit("linear") + assert isinstance(out, MyData) + np.testing.assert_allclose(out.ctx["fit_params"][0], [2.0, 1.0], atol=1e-8) + + def test_fit_window_growth_rate(self): + edges = np.linspace(0.0, 1.0, 61) + centers = 0.5 * (edges[:-1] + edges[1:]) + y = 1.0 * np.exp(2 * 0.5 * centers) + d = _make(MyData, [edges], y[:, np.newaxis]) + out = d.fit("exp2", window=True) + assert isinstance(out, MyData) + assert out.ctx["fit_params"][0][1] == pytest.approx(0.5, abs=1e-2) + + def test_differentiate(self): + edges = np.linspace(0.0, 1.0, 17) + centers = 0.5 * (edges[:-1] + edges[1:]) + d = _make(MyData, [edges], (centers**2)[:, np.newaxis]) + out = d.differentiate() + assert isinstance(out, MyData) + + def test_collect(self): + grid = [np.linspace(0.0, 1.0, 5)] + a = _make(MyData, grid, np.full((4, 1), 2.0), time=0.0) + b = _make(MyData, grid, np.full((4, 1), 3.0), time=1.0) + out = api_verbs.collect(a, b) + assert isinstance(out, MyData) + + def test_sort(self): + grid = [np.linspace(0.0, 1.0, 5)] + a = _make(MyData, grid, np.full((4, 1), 2.0)) + a._file_name = "field_10.gkyl" + b = _make(MyData, grid, np.full((4, 1), 3.0)) + b._file_name = "field_2.gkyl" + out = api_verbs.sort(a, b) + assert [d.file_name for d in out] == ["field_2.gkyl", "field_10.gkyl"] + + def test_evaluate(self): + grid = [np.linspace(0.0, 1.0, 5)] + a = _make(MyData, grid, np.full((4, 1), 2.0)) + b = _make(MyData, grid, np.full((4, 1), 3.0)) + out = api_verbs.evaluate("f0 f1 +", a, b) + assert isinstance(out, MyData) + np.testing.assert_allclose(out.get_values(), 5.0) + + @needs_gkeyll + def test_map(self): + from postgkyl.gpython import basis as gpython_basis + + lower, upper, cells = 0.0, 4.0, 4 + node_eta = gpython_basis.node_coords("serendipity", 1, 1)[:, 0] + n2m = gpython_basis.nodal_to_modal_matrix("serendipity", 1, 1) + dz = (upper - lower) / cells + centers = lower + (np.arange(cells) + 0.5) * dz + nodal_z = centers[:, None] + 0.5 * dz * node_eta[None, :] + modal = nodal_z @ n2m.T # exact per-cell modal coeffs of the identity map + + mapping = GDataState() + mapping.ctx.update(basis_type="serendipity", + poly_order=1, + value_form="modal", + cells=np.array([cells], dtype=np.int64)) + mgrid = [np.linspace(lower, upper, cells + 1)] + mapping.push(mgrid, gpython.GkylArray.from_numpy(modal)) + + target = _make(MyData, [np.linspace(lower, upper, 17)], np.zeros((16, 1))) + out = target.map(mapping, space="conf") + assert isinstance(out, MyData) + np.testing.assert_allclose(out.grid[0], target.grid[0], atol=1e-12) + + @needs_gkeyll + def test_mul_div_interpolate_to_modal_nodal_quad_apply_integrate(self): + F1 = os.path.join( + DATA, + "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + a, b = MyData(F1), MyData(F1) + assert isinstance(a.mul(b), MyData) + a2, b2 = MyData(F1), MyData(F1) + assert isinstance(a2.div(b2), MyData) + assert isinstance(MyData(F1).interpolate(), MyData) + assert isinstance(MyData(F1).to_nodal(), MyData) + assert isinstance(MyData(F1).to_nodal().to_modal(), MyData) + assert isinstance(MyData(F1).to_quad(), MyData) + assert isinstance(MyData(F1).apply(np.abs), MyData) + result = MyData(F1).integrate() + assert result is not None + assert isinstance(MyData(F2D_VEC).integrate(0), MyData) + + +# ============================================================ keyword pass-through +class TestKeywordPassthrough: + + def test_fft_psd_kwarg_reaches_the_verb(self): + d = _line(value=1.0, n=16) + full = d.fft(psd=False) + half = d.fft(psd=True) + assert half.values.shape[0] == full.values.shape[0] // 2 + + def test_magsq_coords_kwarg_reaches_the_verb(self): + d = _make(MyData, [np.linspace(0.0, 1.0, 5)], + np.tile([1.0, 2.0, 3.0], (4, 1))) + default = d.magsq() # "0:3" -> 1+4+9 + partial = d.magsq(coords="0:2") # 1+4 + np.testing.assert_allclose(default.get_values().flat[0], 14.0) + np.testing.assert_allclose(partial.get_values().flat[0], 5.0) + + def test_mask_lower_vs_upper_kwarg_reaches_the_verb(self): + d = _make(MyData, [np.linspace(0.0, 1.0, 6)], np.arange(5.0)[:, np.newaxis]) + lower = d.mask(lower=2.0) + upper = d.mask(upper=2.0) + assert lower.get_values().mask[0, 0] and not lower.get_values().mask[-1, 0] + assert upper.get_values().mask[-1, 0] and not upper.get_values().mask[0, 0] + + +# ======================================================== end-to-end chains +@needs_gkeyll +class TestEndToEndChains: + + def test_interpolate_magsq_plot(self): + fig = pg.load(F2D_VEC).interpolate().magsq().plot(no_show=True) + assert fig is not None + + def test_interpolate_select_fft(self): + # fft's output grid is a frequency axis (one entry per value, not a + # nodal N+1 edge array), so it is not directly re-plottable through the + # same render path as the other chains -- exercised on values instead. + out = pg.load(F1D).interpolate().select(comp=0).fft(psd=True) + assert isinstance(out, pg.GData) + assert out.values.shape[0] == pg.load(F1D).interpolate().select( + comp=0).num_cells[0] // 2 + + def test_interpolate_select_mask_fit(self): + out = pg.load(F1D).interpolate().select(comp=0).mask( + lower=-1e30).fit("linear") + assert isinstance(out, pg.GData) + assert "fit_params" in out.ctx + + +# ================================================================== group +class TestGDataGroup: + + def _frames(self, cls=MyData): + grid = [np.linspace(0.0, 1.0, 5)] + return [ + _make(cls, grid, np.full((4, 1), v), time=t) + for t, v in ((0.0, 1.0), (1.0, 2.0), (2.0, 3.0)) + ] + + def test_broadcast_non_terminal_verb_returns_a_group_of_the_same_class(self): + g = ApiGDataGroup(self._frames()) + out = g.select(comp=0) + assert isinstance(out, ApiGDataGroup) + assert len(out) == 3 + for member in out: + assert isinstance(member, MyData) + + def test_broadcast_chains(self): + g = ApiGDataGroup(self._frames()) + out = g.select(comp=0).mask(lower=-1e30) + assert isinstance(out, ApiGDataGroup) + assert len(out) == 3 + + def test_sort_reorders_members_and_preserves_the_group(self): + frames = self._frames() + names = ("field_10.gkyl", "field_1.gkyl", "field_2.gkyl") + for frame, name in zip(frames, names): + frame._file_name = name + + out = ApiGDataGroup(frames).sort() + + assert isinstance(out, ApiGDataGroup) + assert [member.file_name for member in out + ] == ["field_1.gkyl", "field_2.gkyl", "field_10.gkyl"] + assert isinstance(out.select(comp=0), ApiGDataGroup) + + def test_sort_accepts_reverse(self): + frames = self._frames() + names = ("field_10.gkyl", "field_1.gkyl", "field_2.gkyl") + for frame, name in zip(frames, names): + frame._file_name = name + + out = ApiGDataGroup(frames).sort(reverse=True) + + assert [member.file_name for member in out + ] == ["field_10.gkyl", "field_2.gkyl", "field_1.gkyl"] + + def test_broadcast_terminal_verb_returns_a_plain_list(self): + g = ApiGDataGroup(self._frames()) + inputs = g.extract_input() + assert isinstance(inputs, list) + assert len(inputs) == 3 + + def test_plot_is_not_broadcast_but_one_shared_figure(self): + # A group is a set of datasets that belong together -- above all a + # multiblock family -- so .plot() is ONE figure with every member drawn + # onto it, not one figure per member. + import matplotlib.figure + + g = ApiGDataGroup(self._frames()) + fig = g.plot(no_show=True) + assert isinstance(fig, matplotlib.figure.Figure) + + def test_broadcast_write_returns_a_list_of_paths(self, tmp_path): + g = ApiGDataGroup(self._frames()) + paths = g.save(out_name=str(tmp_path / "frame")) + assert isinstance(paths, list) + assert len(paths) == 3 + for p in paths: + assert os.path.isfile(p) + + def test_broadcast_non_callable_property_returns_a_plain_list(self): + g = ApiGDataGroup(self._frames()) + dims = g.num_dims + assert isinstance(dims, list) + assert len(dims) == 3 + assert all(d == g[0].num_dims for d in dims) + + def test_info_is_explicit_not_broadcast_and_enumerates_members(self): + g = ApiGDataGroup(self._frames()) + summaries = g.info() + assert isinstance(summaries, list) + assert len(summaries) == 3 + assert "#0" in summaries[0] and "#1" in summaries[1] and "#2" in summaries[2] + + def test_collect_combines_members_into_one_dataset(self): + g = ApiGDataGroup(self._frames()) + out = g.collect() + assert isinstance(out, MyData) + np.testing.assert_allclose(out.get_grid()[0], [0.0, 1.0, 2.0]) + + def test_evaluate_combines_named_members(self): + g = ApiGDataGroup(self._frames()[:2]) + out = g.evaluate("f0 f1 +") + assert isinstance(out, MyData) + np.testing.assert_allclose(out.get_values(), 3.0) # 1.0 + 2.0 + + def test_plotly_animate_delegates_the_whole_group(self, monkeypatch): + frames = self._frames() + sentinel = object() + + def fake_plotly_animate(datasets, **kwargs): + assert list(datasets) == frames + assert kwargs == {"fps": 7} + return sentinel + + monkeypatch.setattr(api_verbs, "plotly_animate", fake_plotly_animate) + assert ApiGDataGroup(frames).plotly_animate(fps=7) is sentinel + + @needs_gkeyll + def test_animate_is_explicit_not_broadcast(self): + from matplotlib.animation import FuncAnimation + frames = [pg.load(F1D).interpolate().select(comp=0) for _ in range(3)] + g = ApiGDataGroup(frames) + anim = g.animate(no_show=True) + assert isinstance(anim, FuncAnimation) + + def test_with_and_and_preserve_the_concrete_class(self): + a, b, c = self._frames() + g = ApiGDataGroup([a, b]) + g2 = g.with_(c) + assert isinstance(g2, ApiGDataGroup) + assert len(g2) == 3 + g3 = g & c + assert isinstance(g3, ApiGDataGroup) + + def test_slicing_preserves_the_concrete_class(self): + g = ApiGDataGroup(self._frames()) + sub = g[0:2] + assert isinstance(sub, ApiGDataGroup) + assert len(sub) == 2 + assert isinstance(g[0], MyData) + + def test_private_and_unknown_attributes_are_not_broadcast(self): + g = ApiGDataGroup(self._frames()) + with pytest.raises(AttributeError): + g._not_a_real_attribute + with pytest.raises(AttributeError): + g.this_verb_does_not_exist() + + def test_is_a_core_dataset_group_too(self): + """The fluent group is a genuine subclass of the verb-less container + (mirrors GData/GDataState); every state-reading behavior still holds.""" + g = ApiGDataGroup(self._frames()) + assert isinstance(g, CoreGDataStateGroup) + assert repr(g) == "" + + +# ================================================================== facade +class TestFacade: + + def test_documented_names_resolve(self): + for name in [ + "GData", "load", "GDataGroup", "plot", "info", "integrate", + "interpolate", "select", "represent", "apply", "save", "collect", + "evaluate", "relchange", "animate", "__version__" + ]: + assert hasattr(pg, name), f"postgkyl has no {name!r}" + + def test_all_is_consistent(self): + assert hasattr(pg, "__all__") + for name in pg.__all__: + assert hasattr(pg, name), f"pg.__all__ names {name!r} but it is missing" + + def test_dataset_group_is_the_fluent_one(self): + assert pg.GDataGroup is ApiGDataGroup + + def test_module_verbs_are_the_api_ones(self): + assert pg.collect is api_verbs.collect + assert pg.evaluate is api_verbs.evaluate + assert pg.relchange is api_verbs.relchange + assert pg.animate is render.animate is operations.animate is api_verbs.animate + assert (pg.plotly_animate is render.plotly_animate is + operations.plotly_animate is api_verbs.plotly_animate) + assert pg.sort is api_verbs.sort diff --git a/tests/test_gdatastate_group.py b/tests/test_gdatastate_group.py new file mode 100644 index 00000000..a1c25aad --- /dev/null +++ b/tests/test_gdatastate_group.py @@ -0,0 +1,120 @@ +"""Tests for postgkyl.gdatastate.gdatastategroup.GDataStateGroup -- the verb-less container. + +Ported from tests_bak/test_group.py: only the state-concerned tests survive +(construction, flattening, indexing, iteration, combining, repr). Tests that +exercised broadcasting (``__getattr__`` dispatch to member verbs) or terminal +verbs (``plot``/``info``/``animate``/``plotly_animate``/``collect``/``evaluate``) are +dropped here -- those methods are deferred to the layer-10 fluent group; see +that layer's worklist. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.gdatastate.gdatastategroup import GDataStateGroup +from postgkyl.gdatastate.gdatastate import GDataState + + +def _line(tag: str = "default", offset: float = 0.0) -> GDataState: + d = GDataState(tag=tag) + d.push([np.linspace(0.0, 1.0, 9)], (np.arange(8.0) + offset)[:, None]) + return d + + +class _SubGData(GDataState): + """Stand-in for the fluent ``GData`` subclass (layer 10 adds the real one).""" + + +class TestConstruction: + + def test_from_list(self): + g = GDataStateGroup([_line("a"), _line("b")]) + assert len(g) == 2 + + def test_flattens_nested(self): + g = GDataStateGroup([_line("a"), [_line("b"), _line("c")]]) + assert len(g) == 3 + + def test_flattens_nested_group(self): + inner = GDataStateGroup([_line("b"), _line("c")]) + g = GDataStateGroup([_line("a"), inner]) + assert len(g) == 3 + assert all(isinstance(d, GDataState) for d in g) + + def test_iter_and_index(self): + a, b = _line("a"), _line("b") + g = GDataStateGroup([a, b]) + assert list(g) == [a, b] + assert g[0] is a + + def test_slice_returns_group(self): + g = GDataStateGroup([_line("a"), _line("b"), _line("c")]) + assert isinstance(g[:2], GDataStateGroup) + assert len(g[:2]) == 2 + + def test_rejects_non_gdata(self): + with pytest.raises(TypeError): + GDataStateGroup([1, 2, 3]) + + def test_empty_group_default(self): + g = GDataStateGroup() + assert len(g) == 0 + assert list(g) == [] + + def test_empty_group_from_empty_list(self): + g = GDataStateGroup([]) + assert len(g) == 0 + + def test_group_of_one(self): + a = _line("a") + g = GDataStateGroup([a]) + assert len(g) == 1 + assert g[0] is a + + def test_heterogeneous_member_types(self): + a = _line("a") + b = _SubGData(tag="b") + b.push([np.linspace(0.0, 1.0, 5)], np.arange(4.0)[:, None]) + g = GDataStateGroup([a, b]) + assert len(g) == 2 + assert type(g[0]) is GDataState + assert isinstance(g[1], _SubGData) + + +class TestCombining: + + def test_with_appends(self): + g = GDataStateGroup([_line("a")]).with_(_line("b"), _line("c")) + assert len(g) == 3 + + def test_with_accepts_group(self): + g = GDataStateGroup([_line("a")]).with_(GDataStateGroup([_line("b")])) + assert len(g) == 2 + + def test_and_operator(self): + g = GDataStateGroup([_line("a")]) & GDataStateGroup([_line("b")]) + assert len(g) == 2 + + def test_with_does_not_mutate(self): + g = GDataStateGroup([_line("a")]) + g.with_(_line("b")) + assert len(g) == 1 + + +class TestSequenceAndRepr: + + def test_datasets_is_defensive_copy(self): + a, b = _line("a"), _line("b") + g = GDataStateGroup([a, b]) + members = g.datasets + members.append(_line("c")) + assert len(g) == 2 + + def test_repr_shows_count(self): + g = GDataStateGroup([_line("a"), _line("b")]) + assert repr(g) == "" + + def test_repr_empty(self): + assert repr(GDataStateGroup()) == "" diff --git a/tests/test_gk_fetch_funcs.py b/tests/test_gk_fetch_funcs.py deleted file mode 100644 index 57c1cbd2..00000000 --- a/tests/test_gk_fetch_funcs.py +++ /dev/null @@ -1,571 +0,0 @@ -"""Postgkyl module for testing the gk_quantities fetch functions numerically. - -``test_gk_load_quantity`` drives every registered quantity end to end, but it -only asserts that a dataset comes out the other side: an algebra error in a -fetch function would pass it silently. This module closes that gap by checking -the fetch functions against a case whose moments are known in closed form, a -*shifted Maxwellian* with density n, parallel drift u and temperature T (mass m, -two perpendicular velocity dimensions):: - - M0 = n M2par = n*(u^2 + T/m) - M1 = n*u M2perp = n*(2T/m) - M3par = n*(u^3 + 3*u*T/m) - M3perp = n*u*(2T/m) - -The sharpest check available is that a Maxwellian carries *no* heat flux in the -fluid frame, so ``qpar_fluid``/``qperp_fluid`` must cancel to round-off. Because -"is zero" is also what a badly broken function returns, each vanishing check is -paired with a perturbed case that must come out nonzero and equal to a known -value. - -The DG fields here are constant within each cell, which makes every weak DG -operation (multiply, invert) exact, so the expected values are matched to -near-machine precision rather than to a loose tolerance. -""" -import numpy as np -import pytest - -import postgkyl.utils.gk_quantities.fetch_funcs as ff -import postgkyl.utils.gkeyll_const as gkc -from postgkyl.data import GData - -# Synthetic DG dataset parameters: 1D, p1 serendipity (num_basis = 2). -_POLY_ORDER = 1 -_BASIS_TYPE = "serendipity" -_NUM_BASIS = 2 -_NUM_CELLS = 4 -_NUM_DIMS = 1 - -# Value of the 0th (constant) modal basis function: the cell average of a DG -# field is its 0th coefficient times _PSI0. -_PSI0 = 2.0**(-0.5*_NUM_DIMS) - -# Shifted-Maxwellian parameters. Deliberately not round numbers, and of -# realistic magnitude, so that a wrong formula cannot coincidentally agree. -_MASS = 3.343e-27 # Deuterium. -_DENS = 2.7e19 -_UPAR = 1.3e4 -_TEMP = 9.5e-18 -_VT_SQ = _TEMP/_MASS # Thermal speed squared, T/m. - -# Probe whether the gkylsoft DG-operator library is available; fetch functions -# that use weak multiply/invert are skipped if it is not. -try: - from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops - GkeyllDGops() - _DGOPS_AVAILABLE = True -except Exception: # noqa: BLE001 - any failure means the lib is unusable here - _DGOPS_AVAILABLE = False - -_needs_dgops = pytest.mark.skipif( - not _DGOPS_AVAILABLE, reason="requires the gkylsoft DG library") - - -def _const_gdata(comp_avgs, mass: float = _MASS, charge: float = 1.0) -> GData: - """Return a DG field that is constant in space, with the given cell averages. - - ``comp_avgs`` is a scalar for a single-component field, or a sequence of one - cell average per physical component. Only the 0th modal coefficient of each - component is nonzero, which makes the weak DG operations exact. - """ - avgs = np.atleast_1d(np.asarray(comp_avgs, dtype=float)) - values = np.zeros((_NUM_CELLS, _NUM_BASIS*avgs.size)) - for comp, avg in enumerate(avgs): - values[:, comp*_NUM_BASIS] = avg/_PSI0 - - gdata = GData(ctx={"poly_order": _POLY_ORDER, "basis_type": _BASIS_TYPE, - "mass": mass, "charge": charge}) - gdata.push([np.linspace(0.0, 1.0, _NUM_CELLS + 1)], values) - return gdata - - -def _cell_avg(gdata: GData, comp: int = 0) -> np.ndarray: - """Cell averages of the comp-th physical component of a DG field.""" - return gdata.get_values()[:, comp*_NUM_BASIS]*_PSI0 - - -# Moments of the shifted Maxwellian, as single-component DG fields. -def _m0(): - return _const_gdata(_DENS) - -def _m1(): - return _const_gdata(_DENS*_UPAR) - -def _m2par(): - return _const_gdata(_DENS*(_UPAR**2 + _VT_SQ)) - -def _m2perp(): - return _const_gdata(_DENS*2.0*_VT_SQ) - -def _m3par(): - return _const_gdata(_DENS*(_UPAR**3 + 3.0*_UPAR*_VT_SQ)) - -def _m3perp(): - return _const_gdata(_DENS*_UPAR*2.0*_VT_SQ) - -def _temp(): - return _const_gdata(_TEMP) - - -def _strip_ctx(gdata: GData, *keys) -> GData: - """Drop attributes from a GData's context, as if the file did not carry them.""" - for key in keys: - gdata.ctx.pop(key, None) - return gdata - - -class TestGetCtxVal: - """Resolution of species attributes: '--extra' first, then the file context.""" - - def test_extra_overrides_the_context(self): - """An explicit --extra must win over the attribute stored in the file.""" - gdata = _const_gdata(1.0, mass=_MASS) - assert ff._get_ctx_val(gdata, "mass", mass=999.0) == 999.0 - - def test_extra_array_overrides_the_context_per_species(self): - """The override must hold for per-species arrays too, entry by entry.""" - gdata = _const_gdata(1.0, mass=_MASS) - for species_idx, expected in enumerate([1.0, 2.0, 3.0]): - got = ff._get_ctx_val(gdata, "mass", mass=[1.0, 2.0, 3.0], species_idx=species_idx) - assert got == expected - - def test_context_is_used_when_extra_does_not_carry_the_key(self): - """Without an --extra the file's own attribute is still what is used.""" - gdata = _const_gdata(1.0, mass=_MASS) - assert ff._get_ctx_val(gdata, "mass") == _MASS - assert ff._get_ctx_val(gdata, "mass", charge=999.0) == _MASS - - def test_scalar_extra_is_used_when_the_context_lacks_it(self): - gdata = _strip_ctx(_const_gdata(1.0), "mass") - assert ff._get_ctx_val(gdata, "mass", mass=7.0) == 7.0 - - def test_none_in_context_falls_back_to_extra(self): - gdata = _const_gdata(1.0) - gdata.ctx["mass"] = None - assert ff._get_ctx_val(gdata, "mass", mass=7.0) == 7.0 - - def test_a_scalar_extra_applies_to_every_species(self): - gdata = _strip_ctx(_const_gdata(1.0), "mass") - for species_idx in range(3): - assert ff._get_ctx_val(gdata, "mass", mass=7.0, species_idx=species_idx) == 7.0 - - def test_per_species_array_is_picked_by_species_index(self): - """'--extra mass=1,2,3' must give each species its own value.""" - gdata = _strip_ctx(_const_gdata(1.0), "mass") - for species_idx, expected in enumerate([1.0, 2.0, 3.0]): - got = ff._get_ctx_val(gdata, "mass", mass=[1.0, 2.0, 3.0], species_idx=species_idx) - assert got == expected - - def test_array_without_a_species_index_is_an_error(self): - """An array is meaningless where there is no species to index with.""" - gdata = _strip_ctx(_const_gdata(1.0), "mass") - with pytest.raises(KeyError, match="not computed per species"): - ff._get_ctx_val(gdata, "mass", mass=[1.0, 2.0]) - - def test_too_short_an_array_is_an_error(self): - """Fewer values than species must be reported, not silently wrap around.""" - gdata = _strip_ctx(_const_gdata(1.0), "mass") - with pytest.raises(ValueError, match="only 2 values"): - ff._get_ctx_val(gdata, "mass", mass=[1.0, 2.0], species_idx=2, species="ion2") - - def test_missing_everywhere_is_an_error(self): - gdata = _strip_ctx(_const_gdata(1.0), "mass") - with pytest.raises(KeyError, match="mass"): - ff._get_ctx_val(gdata, "mass") - - -class TestMoments: - """Primitive quantities recovered from the Maxwellian's raw moments.""" - - @_needs_dgops - def test_upar_from_M0_M1(self): - """upar = M1/M0 must return the drift speed the moments were built with.""" - upar = ff.fetch_s1c0_div_s0c0([_m0(), _m1()]) - assert np.allclose(_cell_avg(upar), _UPAR, rtol=1e-12) - - @_needs_dgops - def test_Tpar_from_M0_M1_M2par(self): - """Tpar = m*(M2par - upar*M1)/M0 must return T for a Maxwellian.""" - Tpar = ff.fetch_Tpar_from_M0_M1_M2par([_m0(), _m1(), _m2par()]) - assert np.allclose(_cell_avg(Tpar), _TEMP, rtol=1e-10) - - @_needs_dgops - def test_Tperp_from_M0_M2perp(self): - """Tperp = m*M2perp/(2*M0) must return T for a Maxwellian.""" - Tperp = ff.fetch_Tperp_from_M0_M2perp([_m0(), _m2perp()]) - assert np.allclose(_cell_avg(Tperp), _TEMP, rtol=1e-12) - - def test_temp_from_Tpar_Tperp(self): - """An isotropic (Tpar = Tperp = T) split must average back to T.""" - temp = ff.fetch_temp_from_Tpar_Tperp([_temp(), _temp()]) - assert np.allclose(_cell_avg(temp), _TEMP, rtol=1e-14) - - def test_M2_is_M2par_plus_M2perp(self): - M2 = ff.fetch_s0c0_add_s1c0([_m2par(), _m2perp()]) - expected = _DENS*(_UPAR**2 + _VT_SQ) + _DENS*2.0*_VT_SQ - assert np.allclose(_cell_avg(M2), expected, rtol=1e-12) - - def test_M3_is_M3par_plus_M3perp(self): - M3 = ff.fetch_s0c0_add_s1c0([_m3par(), _m3perp()]) - expected = _DENS*(_UPAR**3 + 3.0*_UPAR*_VT_SQ) + _DENS*_UPAR*2.0*_VT_SQ - assert np.allclose(_cell_avg(M3), expected, rtol=1e-12) - - def test_add_selects_the_requested_components(self): - """fetch_s0c2_add_s0c3 must add components 2 and 3, not whole arrays.""" - gdata = _const_gdata([2.0, 3.0, 4.0, 5.0]) - out = ff.fetch_s0c2_add_s0c3([gdata]) - assert out.get_values().shape[-1] == _NUM_BASIS, "output must be single-component" - assert np.allclose(_cell_avg(out), 4.0 + 5.0, rtol=1e-14) - - @_needs_dgops - def test_press_p(self): - """p = n*T.""" - press = ff.fetch_press_p([_m0(), _temp()]) - assert np.allclose(_cell_avg(press), _DENS*_TEMP, rtol=1e-12) - - -class TestMaxwellianMomentSources: - """Quantities read out of the packed Maxwellian/BiMaxwellian moment files. - - Those files store [n, upar, T/m] and [n, upar, Tpar/m, Tperp/m], so these - tests pin down both the component indexing and the mass normalization. - """ - - def test_Tpar_from_BiMax(self): - bimax = _const_gdata([_DENS, _UPAR, _VT_SQ, _VT_SQ]) - Tpar = ff.fetch_Tpar_from_BiMax([bimax]) - assert np.allclose(_cell_avg(Tpar), _TEMP, rtol=1e-12) - - def test_Tperp_from_BiMax(self): - bimax = _const_gdata([_DENS, _UPAR, _VT_SQ, _VT_SQ]) - Tperp = ff.fetch_Tperp_from_BiMax([bimax]) - assert np.allclose(_cell_avg(Tperp), _TEMP, rtol=1e-12) - - def test_temp_from_Max(self): - maxmom = _const_gdata([_DENS, _UPAR, _VT_SQ]) - temp = ff.fetch_temp_from_Max([maxmom]) - assert np.allclose(_cell_avg(temp), _TEMP, rtol=1e-12) - - @_needs_dgops - def test_press_from_Max(self): - maxmom = _const_gdata([_DENS, _UPAR, _VT_SQ]) - press = ff.fetch_press_from_Max([maxmom]) - assert np.allclose(_cell_avg(press), _DENS*_TEMP, rtol=1e-12) - - @_needs_dgops - def test_press_from_BiMax(self): - bimax = _const_gdata([_DENS, _UPAR, _VT_SQ, _VT_SQ]) - press = ff.fetch_press_from_BiMax([bimax]) - assert np.allclose(_cell_avg(press), _DENS*_TEMP, rtol=1e-12) - - -class TestHeatFluxes: - """Lab-frame energy fluxes and fluid-frame heat fluxes.""" - - def test_qpar_lab_frame(self): - """qpar = (m/2)*M3par.""" - qpar = ff.fetch_qpar([_m3par()]) - expected = 0.5*_MASS*_DENS*(_UPAR**3 + 3.0*_UPAR*_VT_SQ) - assert np.allclose(_cell_avg(qpar), expected, rtol=1e-12) - - def test_qperp_lab_frame(self): - """qperp = (m/2)*M3perp, which is n*u*T for a Maxwellian.""" - qperp = ff.fetch_qperp([_m3perp()]) - assert np.allclose(_cell_avg(qperp), _DENS*_UPAR*_TEMP, rtol=1e-12) - - @_needs_dgops - def test_qpar_fluid_vanishes_for_a_maxwellian(self): - """A Maxwellian carries no parallel heat flux in the fluid frame. - - The three terms of (m/2)*[M3par - 3*u*M2par + 2*u^2*M1] cancel exactly, so - the residual is compared against the size of an individual term rather than - against an absolute zero. - """ - qpar_fluid = ff.fetch_qpar_fluid([_m0(), _m1(), _m2par(), _m3par()]) - term_scale = 0.5*_MASS*_DENS*abs(_UPAR)**3 - assert np.all(np.abs(_cell_avg(qpar_fluid))/term_scale < 1e-10) - - @_needs_dgops - def test_qperp_fluid_vanishes_for_a_maxwellian(self): - qperp_fluid = ff.fetch_qperp_fluid([_m0(), _m1(), _m2perp(), _m3perp()]) - term_scale = 0.5*_MASS*_DENS*abs(_UPAR)*2.0*_VT_SQ - assert np.all(np.abs(_cell_avg(qperp_fluid))/term_scale < 1e-10) - - @_needs_dgops - def test_qpar_fluid_tracks_a_skewed_distribution(self): - """Guard against the vanishing tests passing for a function that is just 0. - - Skewing M3par away from its Maxwellian value by dM3 is a pure heat-flux - perturbation, so the fluid-frame flux must become exactly (m/2)*dM3. - """ - delta_m3 = 0.05*_DENS*_UPAR**3 - m3par_skewed = _const_gdata(_DENS*(_UPAR**3 + 3.0*_UPAR*_VT_SQ) + delta_m3) - - qpar_fluid = ff.fetch_qpar_fluid([_m0(), _m1(), _m2par(), m3par_skewed]) - assert np.allclose(_cell_avg(qpar_fluid), 0.5*_MASS*delta_m3, rtol=1e-8) - - @_needs_dgops - def test_qperp_fluid_tracks_a_skewed_distribution(self): - delta_m3 = 0.05*_DENS*_UPAR*2.0*_VT_SQ - m3perp_skewed = _const_gdata(_DENS*_UPAR*2.0*_VT_SQ + delta_m3) - - qperp_fluid = ff.fetch_qperp_fluid([_m0(), _m1(), _m2perp(), m3perp_skewed]) - assert np.allclose(_cell_avg(qperp_fluid), 0.5*_MASS*delta_m3, rtol=1e-8) - - -@_needs_dgops -class TestThermalSpeed: - """vth = sqrt(T/m), m being the requested species' own mass.""" - - def test_vt(self): - vt = ff.fetch_vt([_temp()]) - assert np.allclose(_cell_avg(vt), np.sqrt(_VT_SQ), rtol=1e-12) - - def test_vt_uses_the_species_mass_from_ctx(self): - """A different species mass must give a different thermal speed.""" - mass = 100.0*_MASS - vt = ff.fetch_vt([_const_gdata(_TEMP, mass=mass)]) - assert np.allclose(_cell_avg(vt), np.sqrt(_TEMP/mass), rtol=1e-12) - - -# --- Sound speed: a two-ion-species plasma ----------------------------------- -# -# A deuterium species (Z=1) and a doubly-charged impurity (Z=2), with the -# electron density set by quasineutrality. Every value is distinct so that a -# formula which mixes up a species, a charge state or a mass cannot accidentally -# agree. -_E_CHARGE = gkc.GKYL_ELEMENTARY_CHARGE - -_N_I1, _T_I1, _M_I1, _Z_I1 = 2.7e19, 6.1e-18, 3.343e-27, 1.0 -_N_I2, _T_I2, _M_I2, _Z_I2 = 4.0e18, 4.3e-18, 2.007e-26, 2.0 -_N_E = _N_I1*_Z_I1 + _N_I2*_Z_I2 # Quasineutrality. -_T_E = 9.5e-18 -_M_E = gkc.GKYL_ELECTRON_MASS - - -def _species_srcs(dens: float, temp: float, mass: float, charge: float) -> list: - """The [M0, temp] source pair for one species, as fetch_c_s receives it.""" - return [_const_gdata(dens, mass=mass, charge=charge), - _const_gdata(temp, mass=mass, charge=charge)] - - -def _elc_srcs(): - return _species_srcs(_N_E, _T_E, _M_E, -_E_CHARGE) - -def _ion1_srcs(): - return _species_srcs(_N_I1, _T_I1, _M_I1, _Z_I1*_E_CHARGE) - -def _ion2_srcs(): - return _species_srcs(_N_I2, _T_I2, _M_I2, _Z_I2*_E_CHARGE) - - -@_needs_dgops -class TestSoundSpeed: - """The multi-species sound speeds, dispatched by '--extra kind='.""" - - def test_ion_acoustic_single_ion_species(self): - """With one Z=1 ion species the formula collapses to sqrt(Te/mi).""" - c_s = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="ion_acoustic") - assert np.allclose(_cell_avg(c_s), np.sqrt(_T_E/_M_I1), rtol=1e-10) - - def test_ion_acoustic_two_ion_species(self): - """c_s = sqrt(Te*sum(n_j*Z_j^2/m_j)/sum(n_j*Z_j)).""" - c_s = ff.fetch_c_s([_elc_srcs(), _ion1_srcs(), _ion2_srcs()], - species=["elc", "ion1", "ion2"], kind="ion_acoustic") - - numer = _N_I1*_Z_I1**2/_M_I1 + _N_I2*_Z_I2**2/_M_I2 - denom = _N_I1*_Z_I1 + _N_I2*_Z_I2 - assert np.allclose(_cell_avg(c_s), np.sqrt(_T_E*numer/denom), rtol=1e-10) - - def test_thermo_single_ion_species(self): - """With one Z=1 ion species: sqrt((gamma_e*Te + gamma_i*Ti)/mi).""" - c_s = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="thermo") - - # n_e = n_i1 here only if quasineutrality holds for a single species, so - # use the general formula rather than the reduced one. - numer = 1.0*_N_E*_T_E + 3.0*_N_I1*_T_I1 - denom = _N_I1*_M_I1 - assert np.allclose(_cell_avg(c_s), np.sqrt(numer/denom), rtol=1e-10) - - def test_thermo_two_ion_species(self): - """c_s = sqrt((gamma_e*n_e*Te + sum(gamma_j*n_j*Tj))/sum(n_j*m_j)).""" - c_s = ff.fetch_c_s([_elc_srcs(), _ion1_srcs(), _ion2_srcs()], - species=["elc", "ion1", "ion2"], kind="thermo") - - numer = 1.0*_N_E*_T_E + 3.0*(_N_I1*_T_I1 + _N_I2*_T_I2) - denom = _N_I1*_M_I1 + _N_I2*_M_I2 - assert np.allclose(_cell_avg(c_s), np.sqrt(numer/denom), rtol=1e-10) - - def test_thermo_defaults_are_gamma_e_1_gamma_i_3(self): - """The documented defaults must be what an un-flagged call actually uses.""" - default = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="thermo") - explicit = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="thermo", - gamma_e=1.0, gamma_i=3.0) - assert np.allclose(_cell_avg(default), _cell_avg(explicit), rtol=1e-12) - - def test_thermo_honours_the_gamma_overrides(self): - c_s = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="thermo", - gamma_e=5.0/3.0, gamma_i=5.0/3.0) - - numer = (5.0/3.0)*(_N_E*_T_E + _N_I1*_T_I1) - assert np.allclose(_cell_avg(c_s), np.sqrt(numer/(_N_I1*_M_I1)), rtol=1e-10) - - def test_default_kind_is_thermo(self): - default = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], species=["elc", "ion1"]) - explicit = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="thermo") - assert np.allclose(_cell_avg(default), _cell_avg(explicit), rtol=1e-12) - - def test_species_order_does_not_matter(self): - """Species are identified by charge sign, so the order is irrelevant.""" - forward = ff.fetch_c_s([_elc_srcs(), _ion1_srcs(), _ion2_srcs()], - species=["elc", "ion1", "ion2"], kind="ion_acoustic") - shuffled = ff.fetch_c_s([_ion2_srcs(), _elc_srcs(), _ion1_srcs()], - species=["ion2", "elc", "ion1"], kind="ion_acoustic") - assert np.allclose(_cell_avg(forward), _cell_avg(shuffled), rtol=1e-12) - - def test_electrons_are_found_by_charge_not_by_name(self): - """A species called anything must still be treated as the electrons.""" - named = ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], - species=["elc", "ion1"], kind="ion_acoustic") - odd = ff.fetch_c_s([_species_srcs(_N_E, _T_E, _M_E, -_E_CHARGE), _ion1_srcs()], - species=["negatron", "deuterium"], kind="ion_acoustic") - assert np.allclose(_cell_avg(named), _cell_avg(odd), rtol=1e-12) - - def test_no_electron_species_is_an_error(self): - with pytest.raises(ValueError, match="exactly one negatively charged"): - ff.fetch_c_s([_ion1_srcs(), _ion2_srcs()], species=["ion1", "ion2"]) - - def test_two_electron_species_is_an_error(self): - with pytest.raises(ValueError, match="exactly one negatively charged"): - ff.fetch_c_s([_elc_srcs(), _elc_srcs(), _ion1_srcs()], - species=["elc1", "elc2", "ion1"]) - - def test_no_ion_species_is_an_error(self): - with pytest.raises(ValueError, match="no positively charged"): - ff.fetch_c_s([_elc_srcs()], species=["elc"]) - - def test_unknown_kind_is_an_error(self): - with pytest.raises(ValueError, match="unknown kind"): - ff.fetch_c_s([_elc_srcs(), _ion1_srcs()], species=["elc", "ion1"], kind="bogus") - - def test_missing_charge_attribute_is_an_error(self): - """Charge missing from both the file and --extra must be reported.""" - srcs = _strip_ctx(_ion1_srcs()[0], "charge"), _ion1_srcs()[1] - with pytest.raises(KeyError, match="charge"): - ff.fetch_c_s([_elc_srcs(), list(srcs)], species=["elc", "ion1"]) - - def test_attributes_can_come_from_per_species_extra_arrays(self): - """Species attributes absent from the files can be given per species. - - This is the '--extra mass=..,..,charge=..,..' path: each species must pick - its own entry, in the order of '--species'. - """ - def bare(dens, temp): - return [_strip_ctx(_const_gdata(dens), "mass", "charge"), - _strip_ctx(_const_gdata(temp), "mass", "charge")] - - c_s = ff.fetch_c_s([bare(_N_E, _T_E), bare(_N_I1, _T_I1), bare(_N_I2, _T_I2)], - species=["elc", "ion1", "ion2"], kind="ion_acoustic", - mass=[_M_E, _M_I1, _M_I2], - charge=[-_E_CHARGE, _Z_I1*_E_CHARGE, _Z_I2*_E_CHARGE]) - - numer = _N_I1*_Z_I1**2/_M_I1 + _N_I2*_Z_I2**2/_M_I2 - denom = _N_I1*_Z_I1 + _N_I2*_Z_I2 - assert np.allclose(_cell_avg(c_s), np.sqrt(_T_E*numer/denom), rtol=1e-10) - - def test_extra_arrays_must_cover_every_species(self): - """Too few values must be reported rather than silently mis-assigned.""" - def bare(dens, temp): - return [_strip_ctx(_const_gdata(dens), "mass", "charge"), - _strip_ctx(_const_gdata(temp), "mass", "charge")] - - with pytest.raises(ValueError, match="only 2 values"): - ff.fetch_c_s([bare(_N_E, _T_E), bare(_N_I1, _T_I1), bare(_N_I2, _T_I2)], - species=["elc", "ion1", "ion2"], - mass=[_M_E, _M_I1, _M_I2], - charge=[-_E_CHARGE, _Z_I1*_E_CHARGE]) - -def _linear_gdata(coeff0: float, coeff1: float) -> GData: - """A single-component p1 field with the given two modal coefficients.""" - values = np.zeros((_NUM_CELLS, _NUM_BASIS)) - values[:, 0] = coeff0 - values[:, 1] = coeff1 - gdata = GData(ctx={"poly_order": _POLY_ORDER, "basis_type": _BASIS_TYPE, - "mass": _MASS, "charge": 1.0}) - gdata.push([np.linspace(0.0, 1.0, _NUM_CELLS + 1)], values) - return gdata - - -def _project_powsqrt_reference(coeff0: float, coeff1: float, exponent: float, - num_quad: int = _POLY_ORDER + 1) -> np.ndarray: - """Independent numpy reference for pow(sqrt(f), exponent) projected on p1 1D. - - Reimplements the quadrature the gkeyll updater performs, from the definition - rather than from its code: the 1D p1 modal basis orthonormal on [-1,1] is - psi0 = 1/sqrt(2), psi1 = sqrt(3/2)*xi, and the projection of g onto it is - coeff_k = integral of g*psi_k over [-1,1], evaluated by Gauss-Legendre. - """ - xi, weights = np.polynomial.legendre.leggauss(num_quad) - psi = np.array([np.full_like(xi, 1.0/np.sqrt(2.0)), np.sqrt(1.5)*xi]) - - f_at_ords = coeff0*psi[0] + coeff1*psi[1] - g_at_ords = np.power(np.sqrt(f_at_ords), exponent) - - return np.array([np.sum(weights*g_at_ords*psi[k]) for k in range(_NUM_BASIS)]) - - -@_needs_dgops -class TestPowSqrt: - """The gkyl_proj_powsqrt_on_basis binding backing vth.""" - - def test_sqrt_of_a_constant_field_is_exact(self): - out = ff._powsqrt_dg(_const_gdata(4.0), 1.0) - assert np.allclose(_cell_avg(out), 2.0, rtol=1e-12) - - def test_sqrt_keeps_the_higher_moments(self): - """A varying field must produce a varying square root. - - This is the whole point of projecting onto the basis rather than taking - the square root of the cell average: the slope must survive. - """ - out = ff._powsqrt_dg(_linear_gdata(4.0/_PSI0, 0.35), 1.0) - assert not np.allclose(out.get_values()[:, 1], 0.0), ( - "sqrt of a varying field must not be piecewise constant") - - def test_sqrt_matches_an_independent_quadrature(self): - """Check the binding against a from-scratch numpy projection.""" - coeff0, coeff1 = 4.0/_PSI0, 0.35 - out = ff._powsqrt_dg(_linear_gdata(coeff0, coeff1),1.0) - - expected = _project_powsqrt_reference(coeff0, coeff1, 1.0) - assert np.allclose(out.get_values()[0, :], expected, rtol=1e-12) - - @pytest.mark.parametrize("exponent", [1.0, -1.0, 3.0]) - def test_exponents_match_an_independent_quadrature(self, exponent): - """sqrt (1), reciprocal sqrt (-1) and the 3/2 power (3).""" - coeff0, coeff1 = 4.0/_PSI0, 0.35 - out = ff._powsqrt_dg(_linear_gdata(coeff0, coeff1), exponent) - - expected = _project_powsqrt_reference(coeff0, coeff1, exponent) - assert np.allclose(out.get_values()[0, :], expected, rtol=1e-12) - - def test_constant_field_exponents(self): - """On a constant field the closed-form answers are exact.""" - field = _const_gdata(4.0) - assert np.allclose(_cell_avg(ff._powsqrt_dg(field, 1.0)), 2.0, rtol=1e-12) - assert np.allclose(_cell_avg(ff._powsqrt_dg(field, -1.0)), 0.5, rtol=1e-12) - assert np.allclose(_cell_avg(ff._powsqrt_dg(field, 3.0)), 8.0, rtol=1e-12) - - def test_multi_component_input_is_rejected(self): - """The kernel has no component index, so a vector field must not be taken.""" - from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops - - field = _const_gdata([1.0, 2.0, 3.0]) # Three physical components. - with pytest.raises(ValueError, match="single-component"): - GkeyllDGops().powsqrt(field, field, 1.0) diff --git a/tests/test_gk_load_quantity.py b/tests/test_gk_load_quantity.py deleted file mode 100644 index 9def466e..00000000 --- a/tests/test_gk_load_quantity.py +++ /dev/null @@ -1,370 +0,0 @@ -"""Postgkyl module for testing the gk-load-quantity command. - -This exercises ``gk_load_quantity`` against *every* quantity registered in -``gk_quant_registry``. Because the on-disk .gkyl writer does not persist the DG -metadata (poly_order, basis_type, mass) needed by the fetch functions, the test -does not rely on real simulation output. Instead it - - 1. creates empty marker files with the exact names each quantity's sources are - discovered by (so ``choose_source`` finds a valid source combination), and - 2. monkeypatches the ``GData`` constructor used inside ``gkquantity`` so that - loading a source returns a small, self-consistent synthetic DG dataset. - -The fetch functions (and their gkylsoft-backed DG operators) then run for real. -Quantities whose computation requires the compiled gkylsoft library are skipped -when that library is unavailable. -""" -import os - -import click -import numpy as np -import pytest - -import postgkyl.commands as cmd -import postgkyl.utils.gk_quantities.gkquantity as gkquantity -from postgkyl.data import GData -from postgkyl.pgkyl import cli -from postgkyl.utils.gk_quantities.registry import gk_quant_registry - -# Synthetic DG dataset parameters: 1D, p1 serendipity (num_basis = 2), four -# physical components so that fetch functions selecting up to component 3 work. -_POLY_ORDER = 1 -_BASIS_TYPE = "serendipity" -_NUM_BASIS = 2 -_NUM_PHYS_COMPS = 4 -_NUM_CELLS = 4 - -# Probe whether the gkylsoft DG-operator library is available. Quantities whose -# fetch functions use it (e.g. press, beta, ExB_vel) are skipped if it is not. -try: - from postgkyl.tools.gkeyll_dg_ops import GkeyllDGops - GkeyllDGops() - _DGOPS_AVAILABLE = True -except Exception: # noqa: BLE001 - any failure means the lib is unusable here - _DGOPS_AVAILABLE = False - -# Extra command options required by specific quantities (beyond the per-component -# selection that every vector quantity needs, which is added automatically below). -_EXTRA_OPTS = {} - -# Species names used in the test. Multi-species quantities (e.g. the sound speed) -# combine an electron species with one or more ion species, so they are requested -# with the whole list; the electron species is identified by its negative charge. -_ELC_SPECIES = "elc" -_ION_SPECIES = "ion" - -# Value of the 0th (constant) modal basis function in 1D: the cell average of a -# DG field is its 0th coefficient times _PSI0. -_PSI0 = 2.0**-0.5 - -# Define the data used to test the handling of GK distribution functions. -_TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), "test_data") -_DISTF_REAL = { - "name": "rt_gk_tcv_iwl_1x2v_p1", - "species": "elc", - "frame": 250, -} - - -def _extra_opts_for(quant) -> str | None: - """Build the '--extra' string a quantity needs to be fetched in the test.""" - opts = [] - if quant.is_vector: - opts.append("dir=0") # Vector quantities need a component selection. - if quant.name in _EXTRA_OPTS: - opts.append(_EXTRA_OPTS[quant.name]) - return ",".join(opts) if opts else None - - -def _make_synthetic_gdata(*args, **kwargs) -> GData: - """Return a small, self-consistent constant-valued DG dataset. - - Each physical component is a positive constant (only the cell-average modal - coefficient is nonzero), which keeps the DG multiply/invert operations - well-defined. Every source is served the same synthetic values, which is - enough to drive the fetch functions; only the charge is read back out of the - file name, so that multi-species quantities (which tell electrons from ions - by the sign of the charge) see a genuine electron species. - """ - values = np.zeros((_NUM_CELLS, _NUM_BASIS * _NUM_PHYS_COMPS)) - for comp in range(_NUM_PHYS_COMPS): - # Distinct positive cell-average per component (1/sqrt(2) is the value of - # the 0th modal serendipity basis function). - values[:, comp * _NUM_BASIS] = (comp + 2) * np.sqrt(2.0) - - file_name = str(args[0]) if args else "" - charge = -1.0 if f"-{_ELC_SPECIES}_" in file_name else 1.0 - - grid = [np.linspace(0.0, 1.0, _NUM_CELLS + 1)] - gdata = GData(ctx={"poly_order": _POLY_ORDER, "basis_type": _BASIS_TYPE, - "mass": 1.0, "charge": charge}) - gdata.push(grid, values) - return gdata - - -def _make_synthetic_gdata_no_attrs(*args, **kwargs) -> GData: - """Synthetic data whose files carry no mass/charge attributes. - - Used to drive the '--extra' fallback, which only kicks in when the attribute - is absent from the file context. - """ - gdata = _make_synthetic_gdata(*args, **kwargs) - gdata.ctx.pop("mass", None) - gdata.ctx.pop("charge", None) - return gdata - - -def _collect_source_files(quant, path: str, name: str, species: str, frame: int) -> set[str]: - """Recursively collect the file names every source combination would look for.""" - files: set[str] = set() - for combo in quant.source: - for src in combo: - if isinstance(src, str): - files.add(quant._src_file_name(path, name, species, src, frame)) - else: - files |= _collect_source_files(src, path, name, species, frame) - return files - - -class TestGkLoadQuantity: - """Test that gk-load-quantity can load every registered quantity.""" - - name = "gktest" - species = "ion" - frame = 0 - - def _make_ctx(self): - ctx = click.core.Context(cli) - ctx.obj = {"data": cmd.DataSpace(), "verbose": False} - return ctx - - @pytest.mark.parametrize("quantity", gk_quant_registry.list()) - def test_load_quantity(self, quantity, tmp_path, monkeypatch): - if quantity == "distf": - self._check_distf_real() - return - - quant = gk_quant_registry.get(quantity) - path = str(tmp_path) - - # A multi-species quantity needs an electron and an ion species to combine. - species = f"{_ELC_SPECIES},{_ION_SPECIES}" if quant.is_multi_species else self.species - - # Create empty marker files for every source so source discovery succeeds. - for species_name in species.split(","): - for file_name in _collect_source_files(quant, path, self.name, species_name, self.frame): - open(file_name, "w").close() - - # Serve synthetic DG data whenever a source file is "loaded". - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata) - - ctx = self._make_ctx() - try: - ctx.invoke( - cmd.gk_load_quantity, - quantity=quantity, - name=self.name, - species=species, - frame=str(self.frame), - path=path, - extra=_extra_opts_for(quant), - ) - except (RuntimeError, FileNotFoundError, OSError) as err: - if not _DGOPS_AVAILABLE: - pytest.skip(f"'{quantity}' requires the gkylsoft DG library: {err}") - raise - - assert ctx.obj["data"].get_num_datasets() >= 1, ( - f"gk-load-quantity produced no dataset for quantity '{quantity}'") - - def _load(self, ctx, quantity, path, species, extra=None): - """Invoke gk-load-quantity for a quantity, skipping if the DG lib is absent.""" - quant = gk_quant_registry.get(quantity) - for species_name in species.split(","): - for file_name in _collect_source_files(quant, path, self.name, species_name, self.frame): - open(file_name, "w").close() - - try: - ctx.invoke( - cmd.gk_load_quantity, - quantity=quantity, - name=self.name, - species=species, - frame=str(self.frame), - path=path, - extra=extra, - ) - except (RuntimeError, FileNotFoundError, OSError) as err: - if not _DGOPS_AVAILABLE: - pytest.skip(f"'{quantity}' requires the gkylsoft DG library: {err}") - raise - - def test_multi_species_quantity_yields_a_single_dataset(self, tmp_path, monkeypatch): - """A multi-species quantity combines its species into one dataset. - - Per-species quantities produce one dataset per requested species; a - multi-species one (the sound speed) must instead fold them all into a - single dataset, which is the whole reason for the separate fetch path. - """ - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata) - species = f"{_ELC_SPECIES},{_ION_SPECIES}" - - ctx = self._make_ctx() - self._load(ctx, "c_s", str(tmp_path), species) - assert ctx.obj["data"].get_num_datasets() == 1, ( - "the sound speed must combine both species into one dataset") - - # The same two species through a per-species quantity give one dataset each. - ctx = self._make_ctx() - self._load(ctx, "temp", str(tmp_path), species) - assert ctx.obj["data"].get_num_datasets() == 2, ( - "a per-species quantity must still produce one dataset per species") - - def test_per_species_extra_array_reaches_each_species(self, tmp_path, monkeypatch): - """'--extra mass=1,2' must give species #0 mass 1 and species #1 mass 2. - - This drives the whole chain end to end: the command parses the array and - tags each species with its index, and _get_ctx_val picks the entry. temp - from the Maxwellian moments is mass*, so the two datasets must - come out differing by exactly the mass ratio. - """ - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata_no_attrs) - - ctx = self._make_ctx() - self._load(ctx, "temp", str(tmp_path), f"{_ELC_SPECIES},{_ION_SPECIES}", extra="mass=1,2") - - assert ctx.obj["data"].get_num_datasets() == 2 - # Several species are tagged per species, in the order they were requested. - first = ctx.obj["data"].get_dataset(0, tag=f"default_{_ELC_SPECIES}").get_values() - second = ctx.obj["data"].get_dataset(0, tag=f"default_{_ION_SPECIES}").get_values() - assert np.allclose(second, 2.0*first), ( - "the second species must be computed with the second mass of the array") - - @pytest.mark.parametrize("extra", [ - "mass=1,2,dir=0", # Pairs separated by commas, as --extra has always been written. - "mass=1,2 dir=0", # Pairs separated by spaces. - "mass=1,2 dir=0", # Extra whitespace. - ]) - def test_extra_pairs_may_be_separated_by_commas_or_spaces(self, extra, tmp_path, monkeypatch): - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata_no_attrs) - - ctx = self._make_ctx() - self._load(ctx, "temp", str(tmp_path), f"{_ELC_SPECIES},{_ION_SPECIES}", extra=extra) - - assert ctx.obj["data"].get_num_datasets() == 2 - first = ctx.obj["data"].get_dataset(0, tag=f"default_{_ELC_SPECIES}").get_values() - second = ctx.obj["data"].get_dataset(0, tag=f"default_{_ION_SPECIES}").get_values() - assert np.allclose(second, 2.0*first), ( - f"'--extra {extra}' must give the two species masses 1 and 2") - - def test_extra_overrides_the_file_attributes(self, tmp_path, monkeypatch): - """'--extra mass=' must win over the mass stored in the output files. - - The synthetic files carry mass=1, so asking for mass=1,2 must change the - result for both species: temp from the Maxwellian moments is - mass*. - """ - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata) # mass=1 in ctx. - species = f"{_ELC_SPECIES},{_ION_SPECIES}" - - ctx = self._make_ctx() - self._load(ctx, "temp", str(tmp_path), species) - from_files = ctx.obj["data"].get_dataset(0, tag=f"default_{_ION_SPECIES}").get_values() - - ctx = self._make_ctx() - self._load(ctx, "temp", str(tmp_path), species, extra="mass=1,2") - overridden = ctx.obj["data"].get_dataset(0, tag=f"default_{_ION_SPECIES}").get_values() - - assert np.allclose(overridden, 2.0*from_files), ( - "--extra mass= must override the mass attribute stored in the files") - - def test_scalar_extra_applies_to_every_species(self, tmp_path, monkeypatch): - """A single '--extra mass=2' must be shared by every species.""" - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata_no_attrs) - - ctx = self._make_ctx() - self._load(ctx, "temp", str(tmp_path), f"{_ELC_SPECIES},{_ION_SPECIES}", extra="mass=2") - - assert ctx.obj["data"].get_num_datasets() == 2 - first = ctx.obj["data"].get_dataset(0, tag=f"default_{_ELC_SPECIES}").get_values() - second = ctx.obj["data"].get_dataset(0, tag=f"default_{_ION_SPECIES}").get_values() - assert np.allclose(second, first) - - def test_multi_species_extra_array_reaches_nested_sources(self, tmp_path, monkeypatch): - """Every species' nested sources must use that species' own array entry. - - c_s(kind=thermo) needs each species' temperature, and temp from the - Maxwellian moments is mass*. So if the sources of every species - were resolved with the same '--extra mass=' entry, the ion temperature would - be built from the electron mass. Only pinning the expected value catches - that: comparing two runs would not, because the mass also enters the - denominator through a correctly-indexed path. - """ - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata_no_attrs) - - mass_e, mass_i = 1.0, 4.0 - ctx = self._make_ctx() - self._load(ctx, "c_s", str(tmp_path), f"{_ELC_SPECIES},{_ION_SPECIES}", - extra=f"kind=thermo,mass={mass_e},{mass_i},charge=-1,1") - - # The synthetic data gives component c the cell average (c+2), so every - # species has n = 2 (component 0) and temp = mass*4 (component 2). - dens, temp_e, temp_i = 2.0, 4.0*mass_e, 4.0*mass_i - expected = np.sqrt((1.0*dens*temp_e + 3.0*dens*temp_i)/(dens*mass_i)) - - values = ctx.obj["data"].get_dataset(0).get_values() - assert np.isclose(values[0, 0]*_PSI0, expected, rtol=1e-10), ( - "each species' temperature must be built from its own '--extra mass=' entry") - - def test_multi_species_quantity_needs_a_species_list(self, tmp_path, monkeypatch): - """Asking for the sound speed without species must say so, not crash oddly.""" - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata) - - ctx = self._make_ctx() - with pytest.raises(ValueError, match="needs a species list"): - ctx.invoke( - cmd.gk_load_quantity, - quantity="c_s", - name=self.name, - species=None, - frame=str(self.frame), - path=str(tmp_path), - extra=None, - ) - - def test_sound_speed_kinds_differ(self, tmp_path, monkeypatch): - """Both --extra kind= values must run and give genuinely different answers.""" - monkeypatch.setattr(gkquantity, "GData", _make_synthetic_gdata) - species = f"{_ELC_SPECIES},{_ION_SPECIES}" - - values = {} - for kind in ("ion_acoustic", "thermo"): - ctx = self._make_ctx() - self._load(ctx, "c_s", str(tmp_path), species, extra=f"kind={kind}") - assert ctx.obj["data"].get_num_datasets() == 1 - values[kind] = ctx.obj["data"].get_dataset(0).get_values().copy() - - assert not np.allclose(values["ion_acoustic"], values["thermo"]), ( - "the two sound-speed definitions should not coincide for this data") - - def _check_distf_real(self): - """ - Test the distf function with real data present in the test_data directory. - """ - ctx = self._make_ctx() - try: - ctx.invoke( - cmd.gk_load_quantity, - quantity="distf", - name=_DISTF_REAL["name"], - species=_DISTF_REAL["species"], - frame=str(_DISTF_REAL["frame"]), - path=_TEST_DATA_DIR, - ) - except (RuntimeError, FileNotFoundError, OSError) as err: - if not _DGOPS_AVAILABLE: - pytest.skip(f"'distf' requires the gkylsoft DG library: {err}") - raise - - assert ctx.obj["data"].get_num_datasets() >= 1, ( - "gk-load-quantity produced no dataset for quantity 'distf'") diff --git a/tests/test_gpython_array.py b/tests/test_gpython_array.py new file mode 100644 index 00000000..7a81f535 --- /dev/null +++ b/tests/test_gpython_array.py @@ -0,0 +1,149 @@ +"""Tests for ``postgkyl.gpython.array.GkylArray`` -- the capsule-owning array. + +Run: PYTHONPATH=src pytest tests/test_gpython_array.py -v +""" + +import gc +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython # noqa: E402 +from postgkyl.gpython.array import GkylArray # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +pytestmark = needs_gkeyll + + +# --------------------------------------------------------------- construction +def test_alloc_is_zeroed_with_the_requested_shape(): + a = GkylArray.alloc(3, 5) + assert (a.ncomp, a.size) == (3, 5) + assert np.array_equal(a.view(), np.zeros((5, 3))) + + +def test_from_numpy_preserves_values_and_shape(): + values = np.arange(2 * 4, dtype=np.float64).reshape(4, 2) + a = GkylArray.from_numpy(values) + assert (a.ncomp, a.size) == (2, 4) + assert np.array_equal(a.view(), values) + + +def test_from_numpy_copies_non_contiguous_input_correctly(): + base = np.arange(40, dtype=np.float64).reshape(10, 4) + sliced = base[::2] # non-contiguous view + assert not sliced.flags["C_CONTIGUOUS"] + a = GkylArray.from_numpy(sliced) + assert np.array_equal(a.view(), sliced) + + +def test_from_numpy_converts_other_dtypes(): + values = np.arange(6, dtype=np.int32).reshape(3, 2) + a = GkylArray.from_numpy(values) + assert a.view().dtype == np.float64 + assert np.array_equal(a.view(), values.astype(np.float64)) + + +def test_clone_is_a_deep_copy(): + a = GkylArray.from_numpy(np.ones((3, 2))) + b = a.clone() + assert np.array_equal(a.view(), b.view()) + # Mutate through the kernel layer (never the view) to prove independence. + gpython.kernels.scale(a, 0.0) # returns a NEW array; `a` itself is untouched + assert np.array_equal(a.view(), np.ones((3, 2))) + assert np.array_equal(b.view(), np.ones((3, 2))) + + +# ---------------------------------------------------------- invalid construction +def test_alloc_rejects_zero_size(): + with pytest.raises(ValueError, match="positive"): + GkylArray.alloc(2, 0) + + +def test_alloc_rejects_zero_ncomp(): + with pytest.raises(ValueError, match="positive"): + GkylArray.alloc(0, 3) + + +def test_alloc_rejects_negative_args(): + with pytest.raises(ValueError, match="positive"): + GkylArray.alloc(-1, 3) + with pytest.raises(ValueError, match="positive"): + GkylArray.alloc(2, -1) + + +def test_from_numpy_rejects_empty_array(): + with pytest.raises(ValueError, match="empty"): + GkylArray.from_numpy(np.zeros((0, 3))) + + +def test_from_numpy_promotes_0d_to_a_single_cell(): + """`np.ascontiguousarray` upgrades a 0-d scalar to shape (1,) before the + extension ever sees it, so this is a valid single-component, single-cell + array, not the `ndim < 1` refusal (which is defensive/unreachable through + this public constructor -- see the C source comment in _gpythonmodule.c).""" + a = GkylArray.from_numpy(np.array(5.0)) + assert (a.ncomp, a.size) == (1, 1) + assert a.view()[0, 0] == 5.0 + + +# --------------------------------------------------------------- memory safety +def test_view_pins_native_memory_after_source_is_dropped(): + """Regression: a view outlives the Python object that produced it.""" + expected = np.arange(6, dtype=np.float64).reshape(3, 2) + v = GkylArray.from_numpy(expected).view() # array is garbage immediately + gc.collect() + assert np.array_equal(v, expected) + + +def test_view_pins_native_memory_for_alloc_too(): + a = GkylArray.alloc(2, 3) + gpython.kernels.shiftc(a, 7.0, 0) # exercise the array without touching `v` + v = a.view() + del a + gc.collect() + assert np.array_equal(v, np.zeros((3, 2))) # `a` was never mutated in place + + +def test_to_numpy_is_a_by_value_copy(): + a = GkylArray.from_numpy(np.ones((2, 2))) + copy = a.to_numpy() + view = a.view() + assert copy.flags.writeable + assert not view.flags.writeable + copy[0, 0] = 99.0 + assert view[0, 0] == 1.0 # the native buffer is untouched + + +def test_view_is_read_only(): + a = GkylArray.alloc(2, 3) + with pytest.raises(ValueError): + a.view()[0, 0] = 1.0 + + +def test_repeated_alloc_and_release_does_not_leak_or_crash(): + for _ in range(500): + a = GkylArray.alloc(4, 10) + a.view() + del a + gc.collect() + + +def test_view_reshapes_with_explicit_cells(): + a = GkylArray.alloc(2, 6) + shaped = a.view(cells=(2, 3)) + assert shaped.shape == (2, 3, 2) + + +def test_repr_reports_shape(): + a = GkylArray.alloc(3, 5) + assert "5 cells" in repr(a) + assert "3 comps" in repr(a) diff --git a/tests/test_gpython_basis.py b/tests/test_gpython_basis.py new file mode 100644 index 00000000..f949e67f --- /dev/null +++ b/tests/test_gpython_basis.py @@ -0,0 +1,264 @@ +"""Tests for ``postgkyl.gpython.basis`` -- Gkeyll basis objects + matrices. + +Run: PYTHONPATH=src pytest tests/test_gpython_basis.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython # noqa: E402 +from postgkyl.gpython import basis as fb # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +pytestmark = needs_gkeyll + + +def _analytic_num_basis(basis_type: str, ndim: int, poly_order: int) -> int: + """Independent (from-scratch) count, NOT derived from the shim's table.""" + if basis_type == "tensor": + return (poly_order + 1)**ndim + # Serendipity: the standard tensor-product-hypercube serendipity finite + # element counts (Arnold & Awanou 2011); 1D collapses to the full + # polynomial space p+1, and 2D matches the textbook 4/8/12-node quad + # elements (bilinear / quadratic-without-center / cubic serendipity). + if ndim == 1: + return poly_order + 1 + if ndim == 2: + return {0: 1, 1: 4, 2: 8, 3: 12}[poly_order] + raise NotImplementedError("no independent closed form wired up for this case") + + +@pytest.mark.parametrize("basis_type,ndim,poly_order", [ + ("serendipity", 1, 0), + ("serendipity", 1, 1), + ("serendipity", 1, 2), + ("serendipity", 1, 3), + ("serendipity", 2, 0), + ("serendipity", 2, 1), + ("serendipity", 2, 2), + ("serendipity", 2, 3), + ("tensor", 1, 2), + ("tensor", 2, 2), + ("tensor", 3, 1), +]) +def test_num_basis_matches_independent_formula(basis_type, ndim, poly_order): + got = fb.num_basis(basis_type, ndim, poly_order) + assert got == _analytic_num_basis(basis_type, ndim, poly_order) + + +def test_analytic_num_basis_helper_rejects_serendipity_3d(): + """The independent reference formula only has closed forms for 1-D/2-D + serendipity; the parametrized cases above never reach 3-D, so this checks + the helper's own guard directly.""" + with pytest.raises(NotImplementedError, match="no independent closed form"): + _analytic_num_basis("serendipity", 3, 1) + + +def test_get_basis_caches_the_same_object(): + a = fb.get_basis("serendipity", 2, 1) + b = fb.get_basis("serendipity", 2, 1) + assert a is b + # Case-insensitivity shares the same cache entry. + c = fb.get_basis("SERENDIPITY", 2, 1) + assert a is c + + +def test_basis_repr(): + b = fb.get_basis("serendipity", 1, 1) + r = repr(b) + assert "serendipity" in r and "ndim=1" in r and "p=1" in r and "N=2" in r + + +# --------------------------------------------------------- boundary guards +@pytest.mark.parametrize( + "basis_type,ndim,poly_order", + [ + ("serendipity", 7, 1), # ndim above Gkeyll's cart_modal_serendip cap + ("serendipity", 0, 1), + ("serendipity", -1, 1), + ("serendipity", 1, 4), # poly_order above the ev[4] table + ("serendipity", 5, 3), # 5D serendipity tops out at p2 + ("serendipity", 6, 2), # 6D serendipity tops out at p1 + ("tensor", 3, 3), # 3D tensor tops out at p2 + ("tensor", 8, 1), + ("bogus", 1, 1), + ]) +def test_unsupported_combinations_raise_cleanly(basis_type, ndim, poly_order): + """These would abort the process or read out-of-bounds C tables if the + Python-side guard were missing (see basis.py's _MAX_POLY_ORDER comment) -- + a clean ValueError, not a crash, is exactly what is being tested here.""" + with pytest.raises(ValueError): + fb.get_basis(basis_type, ndim, poly_order) + + +@pytest.mark.parametrize("basis_type,ndim,poly_order", [ + ("serendipity", 5, 2), + ("serendipity", 6, 1), + ("tensor", 3, 2), +]) +def test_boundary_combinations_that_ARE_supported(basis_type, ndim, poly_order): + b = fb.get_basis(basis_type, ndim, poly_order) + assert (b.ndim, b.poly_order) == (ndim, poly_order) + + +# --------------------------------------------------------------- eval_matrix +def test_eval_matrix_at_cell_center_is_the_constant_mode(): + """b_0 at z=0 is the normalized constant mode 1/sqrt(2)**ndim for + serendipity/tensor (orthonormal on [-1,1]^ndim with respect to dz).""" + for ndim in (1, 2, 3): + m = fb.eval_matrix("serendipity", ndim, 1, np.zeros((1, ndim))) + assert np.isclose(m[0, 0], (1.0 / np.sqrt(2.0))**ndim) + + +def test_eval_matrix_reproduces_an_in_basis_polynomial(): + """Build modal coefficients for f(z) = 1 + 2z + 3z^2 (degree <= p=2) via + nodal_to_modal, then check eval_matrix reproduces f exactly at arbitrary + points (not just the nodes used to build it).""" + basis_type, ndim, p = "serendipity", 1, 2 + nodes = fb.node_coords(basis_type, ndim, p)[:, 0] + + def f(z): + return 1.0 + 2.0 * z + 3.0 * z**2 + + fnodal = f(nodes) + n2m = fb.nodal_to_modal_matrix(basis_type, ndim, p) + coeffs = n2m @ fnodal + + probe = np.linspace(-1, 1, 11).reshape(-1, 1) + m = fb.eval_matrix(basis_type, ndim, p, probe) + got = m @ coeffs + np.testing.assert_allclose(got, f(probe[:, 0]), atol=1e-12) + + +def test_nodal_to_modal_and_modal_to_nodal_are_exact_inverses(): + for basis_type, ndim, p in [("serendipity", 1, 2), ("serendipity", 2, 1), + ("tensor", 2, 2)]: + n2m = fb.nodal_to_modal_matrix(basis_type, ndim, p) + m2n = fb.modal_to_nodal_matrix(basis_type, ndim, p) + nb = fb.num_basis(basis_type, ndim, p) + np.testing.assert_allclose(n2m @ m2n, np.eye(nb), atol=1e-12) + np.testing.assert_allclose(m2n @ n2m, np.eye(nb), atol=1e-12) + + +def test_modal_quad_round_trip_exact_for_in_degree_polynomials(): + """quad_to_modal(modal_to_quad(c)) == c whenever num_quad >= p+1: the + q2m projection integrates b_j(z)*f(z), degree <= 2p, and an n-point + Gauss rule is exact to degree 2n-1, so 2*num_quad-1 >= 2p needs + num_quad >= p+1 (not merely p, as a naive reading of "degree <= p" would + suggest -- this is exactly why the num_quad choice matters here).""" + basis_type, ndim, p, num_quad = "serendipity", 1, 2, 3 + rng = np.random.default_rng(42) + nb = fb.num_basis(basis_type, ndim, p) + coeffs = rng.normal(size=nb) + + m2q = fb.modal_to_quad_matrix(basis_type, ndim, p, num_quad) + q2m = fb.quad_to_modal_matrix(basis_type, ndim, p, num_quad) + back = q2m @ (m2q @ coeffs) + np.testing.assert_allclose(back, coeffs, atol=1e-12) + + +def test_interpolation_matrix_layout_matches_fortran_tensor_order(): + """Row i of a 2D interpolation matrix corresponds to np.unravel_index(i, + [n,n], order='F') -- dimension 0 fastest.""" + n = 3 + pts_1d = fb.interpolation_points_1d(n) + pts_2d = fb.tensor_points(pts_1d, 2) + for i in range(n * n): + idx = np.unravel_index(i, (n, n), order="F") + expected = [pts_1d[idx[0]], pts_1d[idx[1]]] + np.testing.assert_allclose(pts_2d[i], expected) + + +def test_interpolation_matrix_is_cached_and_read_only(): + m1 = fb.interpolation_matrix("serendipity", 1, 1, 2) + m2 = fb.interpolation_matrix("serendipity", 1, 1, 2) + assert m1 is m2 + with pytest.raises(ValueError): + m1[0, 0] = 5.0 + + +def test_gauss_quad_weights_sum_to_domain_volume(): + for ndim in (1, 2, 3): + _, w = fb.gauss_quad(ndim, 3) + assert np.isclose(w.sum(), 2.0**ndim) + + +def test_node_coords_shape(): + coords = fb.node_coords("serendipity", 2, 1) + nb = fb.num_basis("serendipity", 2, 1) + assert coords.shape == (nb, 2) + + +# --------------------------------------------------------------- hybrid/gkhybrid +@pytest.mark.parametrize("basis_type,ndim,expected_num_basis", [ + ("hybrid", 2, 6), + ("hybrid", 3, 12), + ("hybrid", 4, 24), + ("gkhybrid", 2, 6), + ("gkhybrid", 3, 12), + ("gkhybrid", 4, 24), + ("gkhybrid", 5, 48), +]) +def test_hybrid_num_basis_matches_gkeyll_kernel_tables(basis_type, ndim, + expected_num_basis): + """Independent counts from gkeyll's own num_basis_list tables in + gkyl_cart_modal_{hybrid,gkhybrid}_priv.h (and its unit tests + ctest_basis.c), for the (cdim, vdim) split basis.py derives from ndim.""" + b = fb.get_basis(basis_type, ndim, 1) + assert (b.ndim, b.poly_order, b.num_basis, + b.id) == (ndim, 1, expected_num_basis, basis_type) + + +@pytest.mark.parametrize( + "basis_type,ndim,poly_order", + [ + ("hybrid", 1, 1), # below Gkeyll's ndim>1 assert + ("hybrid", 5, 1), # no (cdim, vdim) split Gkeyll compiles kernels for + ("hybrid", 2, 2), # hybrid only exists at poly_order 1 + ("gkhybrid", 1, 1), + ("gkhybrid", 6, 1), + ("gkhybrid", 3, 2), + ]) +def test_hybrid_unsupported_combinations_raise_cleanly(basis_type, ndim, + poly_order): + with pytest.raises(ValueError): + fb.get_basis(basis_type, ndim, poly_order) + + +def test_cdim_vdim_raises_for_unsupported_hybrid_ndim(): + with pytest.raises(ValueError, match="Gkeyll's hybrid basis supports ndim"): + fb.cdim_vdim("hybrid", 5) + + +def test_hybrid_and_gkhybrid_are_distinct_bases_at_the_same_ndim(): + """ndim=2 exists for both; they must not collide in the cache or alias + the same compiled basis.""" + hyb = fb.get_basis("hybrid", 2, 1) + gkhyb = fb.get_basis("gkhybrid", 2, 1) + assert hyb.id == "hybrid" and gkhyb.id == "gkhybrid" + assert hyb is not gkhyb + + +@pytest.mark.parametrize("basis_type,ndim", [ + ("hybrid", 2), + ("hybrid", 3), + ("gkhybrid", 2), + ("gkhybrid", 3), + ("gkhybrid", 4), +]) +def test_hybrid_nodal_to_modal_and_modal_to_nodal_are_exact_inverses( + basis_type, ndim): + n2m = fb.nodal_to_modal_matrix(basis_type, ndim, 1) + m2n = fb.modal_to_nodal_matrix(basis_type, ndim, 1) + nb = fb.num_basis(basis_type, ndim, 1) + np.testing.assert_allclose(n2m @ m2n, np.eye(nb), atol=1e-12) + np.testing.assert_allclose(m2n @ n2m, np.eye(nb), atol=1e-12) diff --git a/tests/test_gpython_kernels.py b/tests/test_gpython_kernels.py new file mode 100644 index 00000000..c845eb24 --- /dev/null +++ b/tests/test_gpython_kernels.py @@ -0,0 +1,709 @@ +"""Tests for ``postgkyl.gpython.kernels`` -- weak algebra, lincomb, reduce, integrate. + +Run: PYTHONPATH=src pytest tests/test_gpython_kernels.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython # noqa: E402 +from postgkyl.gpython import kernels as k # noqa: E402 +from postgkyl.gpython.array import GkylArray # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +pytestmark = needs_gkeyll + + +def _smooth_field(basis_type, ndim, p, cells, rng, shift=0.0): + """Random-but-smooth modal coefficients: only the constant + a small + perturbation on the higher modes, and shifted away from zero so weak + division never divides by (near-)zero.""" + nb = gpython.basis.num_basis(basis_type, ndim, p) + coeffs = rng.normal(scale=0.05, size=(cells, nb)) + coeffs[:, 0] += shift + return GkylArray.from_numpy(coeffs) + + +# --------------------------------------------------------------- weak algebra +@pytest.mark.parametrize("ndim,p", [(1, 1), (1, 2), (2, 1), (2, 2)]) +def test_weak_mul_div_are_inverses_on_smooth_fields(ndim, p): + rng = np.random.default_rng(42) + basis_type = "serendipity" + cells = 6 + a = _smooth_field(basis_type, ndim, p, cells, rng, shift=3.0) + b = _smooth_field(basis_type, ndim, p, cells, rng, shift=5.0) + ab = k.weak_mul(basis_type, ndim, p, a, b) + back = k.weak_div(basis_type, ndim, p, ab, b) + np.testing.assert_allclose(back.view(), a.view(), atol=1e-10) + + +def test_weak_inv_matches_weak_div_by_one(): + rng = np.random.default_rng(7) + basis_type, ndim, p, cells = "serendipity", 1, 1, 4 + a = _smooth_field(basis_type, ndim, p, cells, rng, shift=4.0) + one = GkylArray.from_numpy( + np.zeros((cells, gpython.basis.num_basis(basis_type, ndim, p)))) + # constant field 1: coefficient 0 is 1/normalization, i.e. sqrt(2)**ndim + one.view() # no-op just to document one is unused below (division test) + inv_a = k.weak_inv(basis_type, ndim, p, a) + back = k.weak_mul(basis_type, ndim, p, inv_a, a) + # a * (1/a) == 1: coefficient 0 equals normalization constant, others ~ 0. + expect = np.zeros_like(back.view()) + expect[:, 0] = np.sqrt(2.0) + np.testing.assert_allclose(back.view(), expect, atol=1e-10) + + +def test_weak_mul_rejects_ncomp_not_a_multiple_of_num_basis(): + basis_type, ndim, p = "serendipity", 1, 1 # num_basis == 2 + a = GkylArray.alloc(3, 4) # 3 is not a multiple of 2 + b = GkylArray.alloc(3, 4) + with pytest.raises(ValueError, match="not a multiple"): + k.weak_mul(basis_type, ndim, p, a, b) + + +def test_weak_mul_rejects_shape_mismatch(): + basis_type, ndim, p = "serendipity", 1, 1 + a = GkylArray.alloc(2, 4) + b = GkylArray.alloc(2, 5) # different size + with pytest.raises(ValueError, match="shape mismatch"): + k.weak_mul(basis_type, ndim, p, a, b) + + +def test_weak_ops_reject_unknown_basis_type(): + a = GkylArray.alloc(2, 4) + b = GkylArray.alloc(2, 4) + with pytest.raises(NotImplementedError, match="serendipity"): + k.weak_mul("bogus", 1, 1, a, b) + + +@pytest.mark.parametrize("ndim", [4, 5, 6]) +def test_weak_mul_div_refuse_ndim_above_3(ndim): + """gkyl_dg_bin_ops' kernel tables assert(dim < 4) -- a process abort if + this guard were missing; it must degrade to a clean exception instead.""" + basis = gpython.basis.get_basis("serendipity", ndim, 1) + a = GkylArray.alloc(basis.num_basis, 3) + b = GkylArray.alloc(basis.num_basis, 3) + with pytest.raises(NotImplementedError, match="ndim 1..3"): + k.weak_mul("serendipity", ndim, 1, a, b) + with pytest.raises(NotImplementedError, match="ndim 1..3"): + k.weak_div("serendipity", ndim, 1, a, b) + + +def test_weak_mul_div_refuse_tensor_poly_order_above_table(): + """Tensor mul/div kernels only go to p2 at ndim 2-3 (p3 slot is NULL).""" + a = GkylArray.alloc(16, 3) # shape irrelevant; guard fires first + b = GkylArray.alloc(16, 3) + with pytest.raises(NotImplementedError, match="poly_order 0..2"): + k.weak_mul("tensor", 2, 3, a, b) + + +def test_weak_inv_rejects_non_p1(): + a = GkylArray.alloc(2, 3) + with pytest.raises(NotImplementedError, match="p=1 only"): + k.weak_inv("serendipity", 1, 2, a) + + +@pytest.mark.parametrize("ndim", [4, 5, 6]) +def test_weak_inv_refuses_ndim_above_3(ndim): + """gkyl_dg_inv_op's kernel table has NO bounds check at all for ndim; this + guard is the only thing standing between a call and undefined behavior.""" + basis = gpython.basis.get_basis("serendipity", ndim, 1) + a = GkylArray.alloc(basis.num_basis, 3) + with pytest.raises(NotImplementedError, match="ndim"): + k.weak_inv("serendipity", ndim, 1, a) + + +# --------------------------------------------------- conf-space x phase-space +def test_mul_conf_phase_by_a_unit_constant_conf_field_is_identity_hybrid(): + """Multiplying by a spatially-uniform conf field of true value 1 can never + raise polynomial degree, so it's an EXACT identity on the phase + coefficients regardless of what the weak cross-mul kernel computes -- + this is the 1x1v PKPM pairing (serendipity conf x hybrid phase).""" + cbasis = gpython.basis.get_basis("serendipity", 1, 1) + pbasis = gpython.basis.get_basis("hybrid", 2, 1) + conf_cells, phase_cells = [3], [3, 4] + cop_coeffs = np.zeros((3, cbasis.num_basis)) + cop_coeffs[:, 0] = np.sqrt(2.0) # constant field value 1 (cdim=1) + cop = GkylArray.from_numpy(cop_coeffs) + rng = np.random.default_rng(3) + pop_coeffs = rng.normal(size=(12, pbasis.num_basis)) + pop = GkylArray.from_numpy(pop_coeffs) + out = k.weak_mul_conf_phase("serendipity", 1, "hybrid", 2, 1, conf_cells, + phase_cells, cop, pop) + np.testing.assert_allclose(out.view(), pop_coeffs, atol=1e-10) + + +def test_mul_conf_phase_by_a_unit_constant_conf_field_is_identity_gkhybrid(): + """Same identity check for the 1x2v gyrokinetic pairing (serendipity conf + x gkhybrid phase, cdim=1 vdim=2).""" + cbasis = gpython.basis.get_basis("serendipity", 1, 1) + pbasis = gpython.basis.get_basis("gkhybrid", 3, 1) + conf_cells, phase_cells = [4], [4, 3, 2] + cop_coeffs = np.zeros((4, cbasis.num_basis)) + cop_coeffs[:, 0] = np.sqrt(2.0) + cop = GkylArray.from_numpy(cop_coeffs) + rng = np.random.default_rng(5) + pop_coeffs = rng.normal(size=(24, pbasis.num_basis)) + pop = GkylArray.from_numpy(pop_coeffs) + out = k.weak_mul_conf_phase("serendipity", 1, "gkhybrid", 3, 1, conf_cells, + phase_cells, cop, pop) + np.testing.assert_allclose(out.view(), pop_coeffs, atol=1e-10) + + +def test_mul_conf_phase_by_a_unit_constant_conf_field_is_identity_serendipity(): + """Same-family serendipity conf x serendipity phase also goes through + gkyl_dg_mul_conf_phase_op_range (not the same-basis gkyl_dg_mul_op path, + since cdim != pdim), so it needs its own identity check.""" + cbasis = gpython.basis.get_basis("serendipity", 1, 2) + pbasis = gpython.basis.get_basis("serendipity", 2, 2) + conf_cells, phase_cells = [3], [3, 5] + cop_coeffs = np.zeros((3, cbasis.num_basis)) + cop_coeffs[:, 0] = np.sqrt(2.0) + cop = GkylArray.from_numpy(cop_coeffs) + rng = np.random.default_rng(9) + pop_coeffs = rng.normal(size=(15, pbasis.num_basis)) + pop = GkylArray.from_numpy(pop_coeffs) + out = k.weak_mul_conf_phase("serendipity", 1, "serendipity", 2, 2, conf_cells, + phase_cells, cop, pop) + np.testing.assert_allclose(out.view(), pop_coeffs, atol=1e-10) + + +def test_mul_conf_phase_rejects_ncomp_mismatch(): + cop = GkylArray.alloc(3, 3) # hybrid conf num_basis is 2, not 3 + pop = GkylArray.alloc(6, 12) + with pytest.raises(ValueError, match="single-field only"): + k.weak_mul_conf_phase("serendipity", 1, "hybrid", 2, 1, [3], [3, 4], cop, + pop) + + +def test_mul_conf_phase_rejects_non_serendipity_conf_for_hybrid(): + cop = GkylArray.alloc(2, 3) + pop = GkylArray.alloc(6, 12) + with pytest.raises(NotImplementedError, match="serendipity conf basis"): + k.weak_mul_conf_phase("tensor", 1, "hybrid", 2, 1, [3], [3, 4], cop, pop) + + +def test_mul_conf_phase_rejects_mismatched_ser_ten_families(): + cop = GkylArray.alloc(2, 3) + pop = GkylArray.alloc(4, 15) + with pytest.raises(NotImplementedError, match="phase basis type alone"): + k.weak_mul_conf_phase("tensor", 1, "serendipity", 2, 1, [3], [3, 5], cop, + pop) + + +def test_mul_conf_phase_rejects_kernel_table_gap(): + """pdim=5, cdim=1 has no serendipity cross-mul kernel at all (NULL in + ser_cross_mul_list) -- must raise cleanly, not call through a NULL + function pointer.""" + cop = GkylArray.alloc(2, 2) + pop = GkylArray.alloc(32, 32) + with pytest.raises(NotImplementedError, match="no serendipity conf\\*phase"): + k.weak_mul_conf_phase("serendipity", 1, "serendipity", 5, 1, [2], + [2, 2, 2, 2, 2], cop, pop) + + +def test_mul_conf_phase_rejects_cells_array_size_mismatch(): + cop = GkylArray.alloc(2, 3) + pop = GkylArray.alloc(4, 20) # cells [3, 5] imply size 15, not 20 + with pytest.raises(ValueError, match="incompatible"): + k.weak_mul_conf_phase("serendipity", 1, "serendipity", 2, 1, [3], [3, 5], + cop, pop) + + +def test_mul_conf_phase_rejects_phase_ndim_not_exceeding_conf_ndim(): + cop = GkylArray.alloc(4, 4) + pop = GkylArray.alloc(4, 4) + with pytest.raises(ValueError, match="must exceed"): + k.weak_mul_conf_phase("serendipity", 2, "serendipity", 2, 1, [2, 2], [2, 2], + cop, pop) + + +# ---------------------------------------------------------- coefficient ops +def test_lincomb_matches_numpy(): + rng = np.random.default_rng(1) + a = GkylArray.from_numpy(rng.normal(size=(5, 3))) + b = GkylArray.from_numpy(rng.normal(size=(5, 3))) + out = k.lincomb(2.0, a, -1.5, b) + np.testing.assert_allclose(out.view(), 2.0 * a.view() - 1.5 * b.view()) + + +def test_lincomb_rejects_shape_mismatch(): + a = GkylArray.alloc(2, 4) + b = GkylArray.alloc(3, 4) + with pytest.raises(ValueError, match="shape mismatch"): + k.lincomb(1.0, a, 1.0, b) + + +def test_scale_matches_numpy_and_does_not_mutate_input(): + a = GkylArray.from_numpy(np.arange(6, dtype=np.float64).reshape(3, 2)) + original = a.view().copy() + out = k.scale(a, -2.0) + np.testing.assert_allclose(out.view(), -2.0 * original) + np.testing.assert_allclose(a.view(), original) + + +def test_shiftc_matches_numpy_and_does_not_mutate_input(): + a = GkylArray.from_numpy(np.zeros((3, 2))) + out = k.shiftc(a, 7.0, 1) + expect = np.zeros((3, 2)) + expect[:, 1] = 7.0 + np.testing.assert_allclose(out.view(), expect) + np.testing.assert_allclose(a.view(), np.zeros((3, 2))) + + +# ---------------------------------------------------------------- reductions +def test_reduce_of_constant_coefficients(): + a = GkylArray.from_numpy(np.full((4, 2), 3.0)) + np.testing.assert_allclose(k.reduce(a, k.GKYL_SUM), [12.0, 12.0]) + np.testing.assert_allclose(k.reduce(a, k.GKYL_MIN), [3.0, 3.0]) + np.testing.assert_allclose(k.reduce(a, k.GKYL_MAX), [3.0, 3.0]) + + +def test_dg_reduce_of_constant_field_min_max_match_the_constant(): + """min/max of a truly constant field equal that constant regardless of how + many Gauss-Legendre nodes per cell the kernel evaluates at.""" + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + coeffs = np.zeros((5, nb)) + coeffs[:, 0] = 3.0 * np.sqrt(2.0) # constant mode -> field value 3.0 + a = GkylArray.from_numpy(coeffs) + assert np.isclose(k.dg_reduce(basis_type, ndim, p, a, 0, "min"), 3.0) + assert np.isclose(k.dg_reduce(basis_type, ndim, p, a, 0, "max"), 3.0) + + +def test_dg_reduce_sum_scales_with_cell_count(): + """`sum` totals the per-node field values across every cell (not divided + by node count), so doubling identical cells must exactly double it -- + a cell-count-independent way to check the "sum over the field" semantics + without needing to know the kernel's internal Gauss-node count.""" + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + + def const_field(ncells, value): + coeffs = np.zeros((ncells, nb)) + coeffs[:, 0] = value * np.sqrt(2.0) + return GkylArray.from_numpy(coeffs) + + small = k.dg_reduce(basis_type, ndim, p, const_field(3, 3.0), 0, "sum") + big = k.dg_reduce(basis_type, ndim, p, const_field(6, 3.0), 0, "sum") + assert small > 0 + assert np.isclose(big, 2.0 * small) + + +def test_dg_reduce_min_max_at_the_gauss_legendre_nodes_for_a_linear_field(): + """min/max are evaluated at the basis's Gauss-Legendre quadrature NODES + (interior points), not the cell edges -- so for f(z) = 3 + 2z they equal f + at the nodes nearest each end, not the true f(-1)/f(1) domain extrema. + Serendipity p=1 in 1D uses the 2-point rule at z = +-1/sqrt(3).""" + basis_type, ndim, p = "serendipity", 1, 1 + # modal coefficients of 3 + 2z in the (normalized Legendre) basis: + # b0 = 1/sqrt(2), b1 = sqrt(3/2) z => c0 = 3*sqrt(2), c1 = 2/sqrt(3/2) + c0 = 3.0 * np.sqrt(2.0) + c1 = 2.0 / np.sqrt(1.5) + a = GkylArray.from_numpy(np.array([[c0, c1]])) + node = 1.0 / np.sqrt(3.0) + assert np.isclose(k.dg_reduce(basis_type, ndim, p, a, 0, "min"), + 3.0 - 2.0 * node) + assert np.isclose(k.dg_reduce(basis_type, ndim, p, a, 0, "max"), + 3.0 + 2.0 * node) + + +def test_dg_reduce_rejects_bad_op_and_bad_comp(): + a = GkylArray.alloc(2, 3) + with pytest.raises(ValueError, match="op"): + k.dg_reduce("serendipity", 1, 1, a, 0, "bogus") + with pytest.raises(ValueError, match="comp"): + k.dg_reduce("serendipity", 1, 1, a, 5, "sum") + + +# ----------------------------------------------------------------- integrate +def test_integrate_constant_field_equals_constant_times_volume(): + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + cells = 4 + coeffs = np.zeros((cells, nb)) + coeffs[:, 0] = 2.0 * np.sqrt(2.0) # constant field value 2.0 + a = GkylArray.from_numpy(coeffs) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([2.0]), + "cells": np.array([cells]) + } + result = k.integrate(grid, basis_type, p, a) + np.testing.assert_allclose(result, [2.0 * 2.0]) # value * volume + + +def test_integrate_abs_and_sq_ops(): + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + coeffs = np.zeros((3, nb)) + coeffs[:, 0] = -2.0 * np.sqrt(2.0) # constant field value -2.0 + a = GkylArray.from_numpy(coeffs) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([3.0]), + "cells": np.array([3]) + } + none = k.integrate(grid, basis_type, p, a, op="none") + absr = k.integrate(grid, basis_type, p, a, op="abs") + sq = k.integrate(grid, basis_type, p, a, op="sq") + np.testing.assert_allclose(none, [-6.0]) + np.testing.assert_allclose(absr, [6.0]) + np.testing.assert_allclose(sq, [12.0]) # (-2)^2 * volume(3) = 12 + + +def test_integrate_factor_scales_the_result(): + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + coeffs = np.zeros((2, nb)) + coeffs[:, 0] = np.sqrt(2.0) + a = GkylArray.from_numpy(coeffs) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([2.0]), + "cells": np.array([2]) + } + result = k.integrate(grid, basis_type, p, a, factor=10.0) + np.testing.assert_allclose(result, [20.0]) + + +def test_integrate_rejects_bad_op(): + a = GkylArray.alloc(2, 2) + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": [2]} + with pytest.raises(ValueError, match="op"): + k.integrate(grid, "serendipity", 1, a, op="bogus") + + +def test_integrate_rejects_unsupported_basis_or_poly_order(): + a = GkylArray.alloc(2, 2) + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": [2]} + with pytest.raises(NotImplementedError): + k.integrate(grid, "tensor", 1, a) + with pytest.raises(NotImplementedError): + k.integrate(grid, "serendipity", 3, a) # p3 unsupported by the kernel set + + +def test_integrate_rejects_ndim_above_3(): + basis = gpython.basis.get_basis("serendipity", 4, 1) + a = GkylArray.alloc(basis.num_basis, 6) + grid = { + "ndim": 4, + "lower": np.zeros(4), + "upper": np.ones(4), + "cells": np.array([1, 1, 1, 6]) + } + with pytest.raises(NotImplementedError, match="ndim 1-3"): + k.integrate(grid, "serendipity", 1, a) + + +def test_integrate_rejects_grid_array_mismatch(): + basis_type, ndim, p = "serendipity", 1, 1 + a = GkylArray.alloc(gpython.basis.num_basis(basis_type, ndim, p), 4) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([5]) + } # 5 != a.size (4) + with pytest.raises(ValueError, match="do not cover"): + k.integrate(grid, basis_type, p, a) + + +# ------------------------------------------------------------------ average +def _const_field(basis_type, ndim, p, cells, value): + nb = gpython.basis.num_basis(basis_type, ndim, p) + b0 = 2.0**(-ndim / 2.0) + coeffs = np.zeros((int(np.prod(cells)), nb)) + coeffs[:, 0] = value / b0 + return GkylArray.from_numpy(coeffs) + + +def test_array_average_partial_reduction_of_constant_field_is_exact(): + """A genuinely partial reduction (some dims survive) is a proper weak + contraction: the surviving 1D field's coefficient 0 is the standard + b0-normalized representation of the (unchanged, since the field is + spatially constant) value -- no basis-dependent surprises here.""" + basis_type, p = "serendipity", 1 + cells = [4, 3] + a = _const_field(basis_type, 2, p, cells, 3.0) + grid = { + "ndim": 2, + "lower": np.array([0.0, 0.0]), + "upper": np.array([2.0, 1.0]), + "cells": np.array(cells) + } + out = k.array_average(grid, + basis_type, + p, + ndim_avg=1, + cells_avg=[cells[0]], + avg_dim=[0, 1], + a=a) + expect = np.zeros((cells[0], gpython.basis.num_basis(basis_type, 1, p))) + expect[:, 0] = 3.0 / (2.0**(-1 / 2.0)) + np.testing.assert_allclose(out.view(), expect, atol=1e-10) + + +def test_array_average_full_reduction_unweighted_writes_a_raw_value(): + """The degenerate (every dim averaged) unweighted kernel + (gkyl_array_average_NxYY_avg) writes a single raw VALUE into + coefficient 0 -- not a b0-normalized coefficient, unlike the partial- + reduction case above. This is exactly the asymmetry + ``dg.modal.average`` corrects for (test_dg_modal_average_* in + test_coverage_leaf.py); this test pins the raw kernel behavior itself.""" + basis_type, p = "serendipity", 1 + cells = [4] + a = _const_field(basis_type, 1, p, cells, 3.0) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([2.0]), + "cells": np.array(cells) + } + out = k.array_average(grid, + basis_type, + p, + ndim_avg=1, + cells_avg=[1], + avg_dim=[1], + a=a) + np.testing.assert_allclose(out.view()[0, 0], 3.0, atol=1e-10) + np.testing.assert_allclose(out.view()[0, 1:], 0.0, atol=1e-10) + + +def test_array_average_full_reduction_weighted_by_a_uniform_weight_matches_integrate( +): + """With ANY weight (even spatially uniform), the kernel performs a real + weak division, so the output IS a properly b0-normalized coefficient -- + matching gkyl_array_integrate / volume for a uniform-weight average.""" + basis_type, p = "serendipity", 1 + cells = [4] + value = 3.0 + a = _const_field(basis_type, 1, p, cells, value) + w = _const_field(basis_type, 1, p, cells, 2.0) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([2.0]), + "cells": np.array(cells) + } + out = k.array_average(grid, + basis_type, + p, + ndim_avg=1, + cells_avg=[1], + avg_dim=[1], + a=a, + weight=w) + b0 = 2.0**(-1 / 2.0) + np.testing.assert_allclose(out.view()[0, 0] * b0, value, atol=1e-10) + + +def test_array_average_rejects_unsupported_basis_or_poly_order(): + a = GkylArray.alloc(2, 4) + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": [4]} + with pytest.raises(NotImplementedError, match="serendipity p1-p2"): + k.array_average(grid, "tensor", 1, 1, [1], [1], a) + with pytest.raises(NotImplementedError, match="serendipity p1-p2"): + k.array_average(grid, "serendipity", 3, 1, [1], [1], a) + + +def test_array_average_rejects_ndim_above_3(): + basis = gpython.basis.get_basis("serendipity", 4, 1) + a = GkylArray.alloc(basis.num_basis, 6) + grid = { + "ndim": 4, + "lower": np.zeros(4), + "upper": np.ones(4), + "cells": np.array([1, 1, 1, 6]) + } + with pytest.raises(NotImplementedError, match="ndim 1-3"): + k.array_average(grid, "serendipity", 1, 1, [1, 1, 1, 6], [1, 0, 0, 0], a) + + +def test_array_average_rejects_ncomp_not_single_field(): + basis_type, p = "serendipity", 1 + a = GkylArray.alloc(4, + 4) # 4 comps: 2 fields of num_basis=2, not single-field + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": [4]} + with pytest.raises(ValueError, match="single-field only"): + k.array_average(grid, basis_type, p, 1, [1], [1], a) + + +def test_array_average_rejects_weight_shape_mismatch(): + basis_type, p = "serendipity", 1 + cells = [4] + a = _const_field(basis_type, 1, p, cells, 3.0) + w = GkylArray.alloc(2, 3) # size 3 != a.size (4) + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": np.array(cells)} + with pytest.raises(ValueError, match="weight"): + k.array_average(grid, basis_type, p, 1, [1], [1], a, weight=w) + + +def test_array_average_rejects_grid_array_mismatch(): + basis_type, p = "serendipity", 1 + a = GkylArray.alloc(gpython.basis.num_basis(basis_type, 1, p), 4) + grid = { + "ndim": 1, + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([5]) + } # 5 != a.size (4) + with pytest.raises(ValueError, match="do not cover"): + k.array_average(grid, basis_type, p, 1, [1], [1], a) + + +# -------------------------------------------------------------- differentiate +def test_weak_differentiate_rejects_poly_order_above_table(): + a = GkylArray.alloc(2, 3) + with pytest.raises(NotImplementedError, match="poly_order 1..2"): + k.weak_differentiate("serendipity", 1, 3, dir=0, diff_order=1, dx=1.0, a=a) + + +def test_weak_differentiate_rejects_dir_out_of_range(): + a = GkylArray.alloc(2, 3) + with pytest.raises(ValueError, match="out of range"): + k.weak_differentiate("serendipity", 1, 1, dir=5, diff_order=1, dx=1.0, a=a) + + +def test_weak_differentiate_rejects_bad_diff_order(): + a = GkylArray.alloc(2, 3) + with pytest.raises(ValueError, match="order must be 1 or 2"): + k.weak_differentiate("serendipity", 1, 1, dir=0, diff_order=3, dx=1.0, a=a) + + +# ----------------------------------------------------------- eval_at_coord_proj +def test_eval_at_coord_proj_rejects_gkhybrid_poly_order_above_1(): + a = GkylArray.alloc(2, 2) + grid = { + "ndim": 3, + "lower": [0.0, 0.0, 0.0], + "upper": [1.0, 1.0, 1.0], + "cells": [1, 1, 1] + } + with pytest.raises(NotImplementedError, match="poly_order 1 only"): + k.eval_at_coord_proj("gkhybrid", 3, 2, 1, grid, [0], [0.0], 1, [1], a) + + +def test_eval_at_coord_proj_rejects_unknown_basis_type(): + a = GkylArray.alloc(2, 2) + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": [1]} + with pytest.raises(NotImplementedError, match="serendipity/tensor/gkhybrid"): + k.eval_at_coord_proj("hybrid", 1, 1, 1, grid, [0], [0.0], 1, [1], a) + + +def test_eval_at_coord_proj_rejects_tensor_ndim_above_table(): + a = GkylArray.alloc(2, 2) + grid = {"ndim": 4, "lower": [0.0] * 4, "upper": [1.0] * 4, "cells": [1] * 4} + with pytest.raises(NotImplementedError, match="ndim"): + k.eval_at_coord_proj("tensor", 4, 1, 4, grid, [0], [0.0], 1, [1], a) + + +def test_eval_at_coord_proj_rejects_poly_order_above_table(): + a = GkylArray.alloc(2, 2) + grid = {"ndim": 1, "lower": [0.0], "upper": [1.0], "cells": [1]} + with pytest.raises(NotImplementedError, match="poly_order 1..2"): + k.eval_at_coord_proj("serendipity", 1, 3, 1, grid, [0], [0.0], 1, [1], a) + + +def test_eval_at_coord_proj_rejects_eval_dirs_out_of_range(): + a = GkylArray.alloc(2, 2) + grid = {"ndim": 2, "lower": [0.0, 0.0], "upper": [1.0, 1.0], "cells": [1, 1]} + with pytest.raises(ValueError, match="out of range"): + k.eval_at_coord_proj("serendipity", 2, 1, 2, grid, [5], [0.0], 1, [1], a) + + +def test_eval_at_coord_proj_rejects_grid_array_mismatch(): + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + a = GkylArray.alloc(nb, 4) + grid = {"ndim": ndim, "lower": [0.0], "upper": [1.0], "cells": [5]} + with pytest.raises(ValueError, match="do not cover"): + k.eval_at_coord_proj(basis_type, ndim, p, ndim, grid, [0], [0.0], 1, [1], a) + + +def test_eval_at_coord_proj_rejects_eval_dirs_coords_length_mismatch(): + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + a = GkylArray.alloc(nb, 4) + grid = {"ndim": ndim, "lower": [0.0], "upper": [1.0], "cells": [4]} + with pytest.raises(ValueError, match="same length"): + k.eval_at_coord_proj(basis_type, ndim, p, ndim, grid, [0], [0.1, 0.2], 1, + [1], a) + + +# ------------------------------------------------------------------ powsqrt +def test_powsqrt_of_a_constant_field_is_exact(): + basis_type, ndim, p, cells = "serendipity", 1, 1, 4 + nb = gpython.basis.num_basis(basis_type, ndim, p) + b0 = 2.0**(-ndim / 2.0) + coeffs = np.zeros((cells, nb)) + coeffs[:, 0] = 4.0 / b0 # constant field value 4.0 + a = GkylArray.from_numpy(coeffs) + out = k.powsqrt(basis_type, ndim, p, [cells], a, 1.0) + np.testing.assert_allclose(out.view()[:, 0] * b0, 2.0, atol=1e-12) + + +@pytest.mark.parametrize("exponent", [1.0, -1.0, 3.0]) +def test_powsqrt_matches_the_apply_pointwise_quadrature_path(exponent): + """``gkyl_proj_powsqrt_on_basis`` and ``dg.rep.apply_pointwise`` both + project through the same modal<->quadrature matrices (``basis.py``'s + Gauss-Legendre rule), so they must agree to quadrature precision on a + genuinely varying (non-constant) field -- this is the cross-check + ``REFACTOR_GKEYLL_FFI.md``'s ``.apply()`` verb already exercises, + independent of the compiled kernel.""" + from postgkyl import dg + + basis_type, ndim, p, cells = "serendipity", 1, 1, 4 + rng = np.random.default_rng(3) + nb = gpython.basis.num_basis(basis_type, ndim, p) + coeffs = rng.normal(scale=0.05, size=(cells, nb)) + coeffs[:, 0] += 4.0 # shifted positive so pow(sqrt(.), .) stays well-defined + a = GkylArray.from_numpy(coeffs) + num_quad = p + 1 + + out = k.powsqrt(basis_type, ndim, p, [cells], a, exponent, num_quad=num_quad) + expect = dg.rep.apply_pointwise( + basis_type, ndim, p, a, + lambda v: np.power(np.sqrt(np.where(v < 0, 1e-40, v)), exponent), + num_quad) + np.testing.assert_allclose(out.view(), expect.view(), atol=1e-10) + + +def test_powsqrt_rejects_multi_component_input(): + """The kernel has no field-index argument, so a multi-component (vector) + field must be refused here (looping per field is ``dg.modal.powsqrt``'s + job, not this thin binding's).""" + basis_type, ndim, p, cells = "serendipity", 1, 1, 4 + nb = gpython.basis.num_basis(basis_type, ndim, p) + a = GkylArray.alloc(3 * nb, cells) # 3 physical components + with pytest.raises(ValueError, match="single-field only"): + k.powsqrt(basis_type, ndim, p, [cells], a, 1.0) + + +def test_powsqrt_rejects_num_quad_below_one(): + basis_type, ndim, p, cells = "serendipity", 1, 1, 4 + nb = gpython.basis.num_basis(basis_type, ndim, p) + a = GkylArray.alloc(nb, cells) + with pytest.raises(ValueError, match="num_quad"): + k.powsqrt(basis_type, ndim, p, [cells], a, 1.0, num_quad=0) + + +def test_powsqrt_rejects_cells_not_covering_the_array(): + basis_type, ndim, p = "serendipity", 1, 1 + nb = gpython.basis.num_basis(basis_type, ndim, p) + a = GkylArray.alloc(nb, 4) + with pytest.raises(ValueError, match="do not cover"): + k.powsqrt(basis_type, ndim, p, [5], a, 1.0) diff --git a/tests/test_gpython_lib.py b/tests/test_gpython_lib.py new file mode 100644 index 00000000..ca804474 --- /dev/null +++ b/tests/test_gpython_lib.py @@ -0,0 +1,157 @@ +"""Tests for ``postgkyl.gpython._lib`` -- the capability-switch handshake. + +Run: PYTHONPATH=src pytest tests/test_gpython_lib.py -v +""" + +import importlib.util +import os +import sys +import types + +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython # noqa: E402 +from postgkyl.gpython import _lib # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +@needs_gkeyll +def test_available_true_when_extension_loaded(): + assert _lib.available() is True + + +@needs_gkeyll +def test_require_returns_the_extension_module(): + mod = _lib.require() + assert mod is sys.modules["postgkyl.gpython._gpython"] + + +@needs_gkeyll +def test_lib_path_points_at_the_loaded_extension(): + p = _lib.lib_path() + assert p is not None + assert p.name.startswith("_gpython") + assert p.exists() + + +@needs_gkeyll +def test_handshake_version_matches(): + g0 = _lib.require() + assert g0.api_version() == g0.GPYTHON_API_VERSION + + +def test_available_false_when_extension_absent(monkeypatch): + """Simulate a no-library install by monkeypatching the module attributes + (the pattern the layer instructions call out explicitly) rather than + reloading the real module in place -- `monkeypatch` guarantees the original + ``_mod``/``_ERROR`` are restored even if an assertion below fails, so this + can never leak a broken capability switch into the rest of the suite.""" + monkeypatch.setattr(_lib, "_mod", None) + monkeypatch.setattr(_lib, "_ERROR", "simulated: no _gpython.so found") + assert _lib.available() is False + with pytest.raises(RuntimeError, match="simulated: no _gpython.so found"): + _lib.require() + assert _lib.lib_path() is None + + +def _exec_independent_lib_copy(): + """Execute a fresh, independent copy of _lib.py's module code. + + Distinct from `postgkyl.gpython._lib` (a different module object entirely) so + mutating its state can never affect `postgkyl.gpython.available`/`require`, + which are bound to the real module's original functions. Its relative + `from . import _gpython` still resolves against the real `postgkyl.gpython` + package, which the caller controls via `sys.modules['postgkyl.gpython._gpython']` + for the duration of the call. + """ + spec = importlib.util.spec_from_file_location( + "postgkyl.gpython._lib_independent_copy", _lib.__file__) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +class _patched_gpython: + """Context manager that makes `from . import _gpython` see `replacement`. + + `from package import submodule` tries `getattr(package, submodule)` + BEFORE consulting `sys.modules`, and the real `postgkyl.gpython` package + object already carries a `_gpython` attribute (set as a side effect of the + real import at process start) -- so patching `sys.modules` alone is not + enough. Both are patched here and restored unconditionally. + """ + + def __init__(self, replacement): + self._replacement = replacement + + def __enter__(self): + self._pkg = sys.modules["postgkyl.gpython"] + self._had_attr = hasattr(self._pkg, "_gpython") + self._old_attr = getattr(self._pkg, "_gpython", None) + self._old_sys_mod = sys.modules.get("postgkyl.gpython._gpython") + if self._had_attr: + delattr(self._pkg, "_gpython") + sys.modules["postgkyl.gpython._gpython"] = self._replacement + + def __exit__(self, *exc): + if self._had_attr: + setattr(self._pkg, "_gpython", self._old_attr) + if self._old_sys_mod is not None: + sys.modules["postgkyl.gpython._gpython"] = self._old_sys_mod + else: + del sys.modules["postgkyl.gpython._gpython"] + return False + + +def test_import_error_when_extension_missing(): + """The actual `try: from . import _gpython / except ImportError` branch.""" + with _patched_gpython(None): # sentinel: forces ImportError + copy = _exec_independent_lib_copy() + + assert copy.available() is False + with pytest.raises(RuntimeError, match="Build the compiled bridge"): + copy.require() + assert copy.lib_path() is None + # The real package's bindings must be entirely unaffected by the above. + assert gpython.available() is True + assert isinstance(gpython.require(), types.ModuleType) + + +@needs_gkeyll +def test_patched_gpython_cleans_up_sys_modules_when_never_previously_imported(): + """``_patched_gpython.__exit__``'s cleanup has two cases: restore whatever was + in ``sys.modules`` before (exercised by every other test here, since the + real ``_gpython`` is always already imported in this environment), or delete + the key entirely when there was nothing to restore. Simulate the latter by + removing the real module first and restoring it manually afterward.""" + real = sys.modules.pop("postgkyl.gpython._gpython") + try: + with _patched_gpython(types.SimpleNamespace()): + assert "postgkyl.gpython._gpython" in sys.modules + assert "postgkyl.gpython._gpython" not in sys.modules + finally: + sys.modules["postgkyl.gpython._gpython"] = real + + +@needs_gkeyll +def test_version_mismatch_degrades_like_missing(): + """A stale `_gpython.so` (wrong GPYTHON_API_VERSION) must degrade the same way.""" + real = sys.modules["postgkyl.gpython._gpython"] + fake = types.SimpleNamespace( + api_version=lambda: real.GPYTHON_API_VERSION + 1000, + GPYTHON_API_VERSION=real.GPYTHON_API_VERSION) + with _patched_gpython(fake): + copy = _exec_independent_lib_copy() + + assert copy.available() is False + with pytest.raises(RuntimeError, match="version mismatch"): + copy.require() + # Unaffected real bindings. + assert gpython.available() is True + assert gpython.require() is real diff --git a/tests/test_gpython_rio.py b/tests/test_gpython_rio.py new file mode 100644 index 00000000..344f984b --- /dev/null +++ b/tests/test_gpython_rio.py @@ -0,0 +1,254 @@ +"""Tests for ``postgkyl.gpython.rio`` -- file I/O through Gkeyll's ``gkyl_array_rio``. + +Run: PYTHONPATH=src pytest tests/test_gpython_rio.py -v +""" + +import glob +import os +import sys +import tempfile + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import gpython # noqa: E402 +from postgkyl.gpython import rio # noqa: E402 +from postgkyl.gpython.array import GkylArray # noqa: E402 +from postgkyl.io.gkyl_reader import GkylReader # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +DATA = os.path.join(ROOT, "tests", "test_data") +FIELD_FILES = sorted( + glob.glob(os.path.join(DATA, "rt_gk_tcv_iwl_1x2v_p1-*.gkyl"))) +# These cross-checks cover field files. Read the independent Python reader's +# header instead of inferring file types from names: dynvectors need not have +# a "_dynvec" suffix (for example, exponential_energy.gkyl). +GENERATED_FILES = [] +for path in sorted(glob.glob(os.path.join(DATA, "generated", "*.gkyl"))): + reader = GkylReader(path, ctx={}) + reader.preload() + if reader.file_type in rio.FIELD_FILE_TYPES: + GENERATED_FILES.append(path) + +pytestmark = needs_gkeyll + +# A non-field (dynvector) file, used below to check that `file_type` correctly +# excludes it from the field-file cross-check. +_NON_FIELD_FILE = None +if gpython.available(): + from postgkyl.gpython import rio as _rio + _dynvec_dir = tempfile.mkdtemp() + _NON_FIELD_FILE = os.path.join(_dynvec_dir, "not_a_field_dynvec.gkyl") + _rio.write_dynvec(_NON_FIELD_FILE, np.array([0.0, 1.0]), + np.array([[1.0], [2.0]])) + + +# ------------------------------------------------------ cross-check vs GkylReader +def _read_with_pure_python(path): + r = GkylReader(path, ctx={}) + r.preload() + return r.load() + + +@pytest.mark.parametrize("path", + FIELD_FILES + GENERATED_FILES, + ids=os.path.basename) +def test_read_field_matches_the_pure_python_reader(path): + """The strongest test in this layer: for every fixture the C reader + accepts, its grid/cells/coefficients must agree exactly with the + independent pure-Python implementation reading the same bytes.""" + py_grid, py_values = _read_with_pure_python(path) + + assert rio.file_type(path) in rio.FIELD_FILE_TYPES + c_grid, c_arr = rio.read_field(path) + c_values = c_arr.to_numpy(cells=c_grid["cells"]) + + assert c_grid["ndim"] == len(py_grid) + for d in range(c_grid["ndim"]): + np.testing.assert_allclose( + py_grid[d], + np.asarray( + np.linspace(c_grid["lower"][d], c_grid["upper"][d], + int(c_grid["cells"][d]) + 1))) + np.testing.assert_allclose(c_values.squeeze(), np.squeeze(py_values)) + + +def test_file_type_of_a_field_file(): + assert rio.file_type(FIELD_FILES[0]) in rio.FIELD_FILE_TYPES + + +def test_file_type_of_a_dynvec_file_is_not_a_field_type(): + assert rio.file_type(_NON_FIELD_FILE) not in rio.FIELD_FILE_TYPES + + +def test_file_type_nonexistent_path_returns_sentinel(): + """`file_type` is documented to return -1 for "not a gkyl file" rather + than raise -- a nonexistent path is exactly that case.""" + assert rio.file_type("/no/such/file.gkyl") == -1 + + +def test_file_type_non_gkyl_file_returns_sentinel(tmp_path): + bogus = tmp_path / "not_a_gkyl_file.gkyl" + bogus.write_bytes(b"definitely not a gkyl binary file") + assert rio.file_type(str(bogus)) == -1 + + +def test_read_header_nonexistent_path_raises(): + with pytest.raises(OSError): + rio.read_header("/no/such/file.gkyl") + + +def test_read_field_nonexistent_path_raises(): + with pytest.raises(OSError): + rio.read_field("/no/such/file.gkyl") + + +def test_read_field_non_gkyl_file_refuses_cleanly(tmp_path): + bogus = tmp_path / "not_a_gkyl_file.gkyl" + bogus.write_bytes(b"this is definitely not a gkyl binary file, at all!!") + with pytest.raises(OSError): + rio.read_field(str(bogus)) + + +def test_read_header_reports_metadata_for_a_modal_file(): + grid, ftype, meta, esznc, tot_cells = rio.read_header(FIELD_FILES[0]) + assert ftype in rio.FIELD_FILE_TYPES + assert esznc > 0 + assert tot_cells == int(np.prod(grid["cells"])) + assert isinstance(meta, + bytes) and len(meta) > 0 # this fixture has msgpack meta + + +# -------------------------------------------------------------------- writing +def test_write_field_round_trips_bit_exactly(tmp_path): + rng = np.random.default_rng(0) + values = rng.normal(size=(4, 3, 2)).astype(np.float64) + arr = GkylArray.from_numpy(values) + grid = { + "lower": np.array([0.0, -1.0]), + "upper": np.array([2.0, 1.0]), + "cells": np.array([4, 3]) + } + path = str(tmp_path / "roundtrip.gkyl") + rio.write_field(path, grid, arr) + + back_grid, back_arr = rio.read_field(path) + np.testing.assert_array_equal(back_grid["lower"], grid["lower"]) + np.testing.assert_array_equal(back_grid["upper"], grid["upper"]) + np.testing.assert_array_equal(back_grid["cells"], grid["cells"]) + np.testing.assert_array_equal(back_arr.to_numpy(), arr.to_numpy()) + + +def test_write_field_with_metadata_round_trips_the_bytes(tmp_path): + import msgpack + meta = msgpack.packb({"polyOrder": 1, "basisType": "serendipity"}) + arr = GkylArray.from_numpy(np.ones((3, 2))) + grid = { + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([3]) + } + path = str(tmp_path / "with_meta.gkyl") + rio.write_field(path, grid, arr, meta=meta) + + _, ftype, back_meta, _, _ = rio.read_header(path) + assert msgpack.unpackb(back_meta) == { + "polyOrder": 1, + "basisType": "serendipity" + } + + +def test_write_field_is_readable_by_the_pure_python_reader(tmp_path): + """Interoperability: a file this floor writes must be a real, standard + .gkyl file, not merely self-consistent with this floor's own reader.""" + arr = GkylArray.from_numpy(np.arange(10, dtype=np.float64).reshape(5, 2)) + grid = { + "lower": np.array([0.0]), + "upper": np.array([5.0]), + "cells": np.array([5]) + } + path = str(tmp_path / "interop.gkyl") + rio.write_field(path, grid, arr) + + py_grid, py_values = _read_with_pure_python(path) + np.testing.assert_allclose(py_grid[0], np.linspace(0.0, 5.0, 6)) + np.testing.assert_allclose(np.squeeze(py_values), arr.to_numpy()) + + +def test_write_field_rejects_grid_array_mismatch(tmp_path): + arr = GkylArray.alloc(2, 4) + grid = { + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([5]) + } + with pytest.raises(ValueError, match="do not cover"): + rio.write_field(str(tmp_path / "bad.gkyl"), grid, arr) + + +def test_write_field_bad_path_raises_oserror(): + arr = GkylArray.alloc(2, 3) + grid = { + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([3]) + } + with pytest.raises(OSError): + rio.write_field("/no/such/directory/out.gkyl", grid, arr) + + +# ------------------------------------------------------------------ dynvector +def test_dynvec_write_read_round_trip(tmp_path): + time = np.array([0.0, 0.1, 0.25, 0.4]) + data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], + [10.0, 11.0, 12.0]]) + path = str(tmp_path / "series.gkyl") + rio.write_dynvec(path, time, data) + + back_time, back_data = rio.read_dynvec(path) + np.testing.assert_allclose(back_time, time) + np.testing.assert_allclose(back_data, data) + + +def test_dynvec_write_read_round_trip_single_component(tmp_path): + time = np.array([0.0, 1.0, 2.0]) + data = np.array([1.5, -2.5, 3.5]) + path = str(tmp_path / "series_1c.gkyl") + rio.write_dynvec(path, time, data) + + back_time, back_data = rio.read_dynvec(path) + np.testing.assert_allclose(back_time, time) + np.testing.assert_allclose(back_data[:, 0], data) + + +def test_dynvec_write_rejects_length_mismatch(tmp_path): + time = np.array([0.0, 1.0, 2.0]) + data = np.array([[1.0], [2.0]]) # only 2 rows + with pytest.raises(ValueError, match="samples"): + rio.write_dynvec(str(tmp_path / "bad.gkyl"), time, data) + + +def test_dynvec_read_nonexistent_file_raises(): + with pytest.raises(OSError): + rio.read_dynvec("/no/such/dynvec.gkyl") + + +def test_dynvec_read_non_dynvec_file_raises(tmp_path): + """A well-formed FIELD file is not a dynvector -- must refuse, not + silently misinterpret the bytes.""" + arr = GkylArray.alloc(2, 3) + grid = { + "lower": np.array([0.0]), + "upper": np.array([1.0]), + "cells": np.array([3]) + } + path = str(tmp_path / "field_not_dynvec.gkyl") + rio.write_field(path, grid, arr) + with pytest.raises(OSError): + rio.read_dynvec(path) diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py deleted file mode 100644 index 67eeefc5..00000000 --- a/tests/test_interpolate.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Postgkyl module for testing DG interpolation""" -import os -import numpy as np - -import postgkyl as pg - - -class TestGkylInterpolate: - """Test Postgkyl interpolate functions.""" - dir_path = f"{os.path.dirname(__file__)}/test_data" - - def test_ser_p1(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - dg = pg.GInterpModal(data, poly_order=1, basis_type="ms") - grid, values = dg.interpolate() - np.testing.assert_equal(len(grid[0]), 17) - np.testing.assert_equal(len(grid[1]), 17) - np.testing.assert_array_equal(values.shape, (16, 16, 1)) - np.testing.assert_approx_equal(values.mean(), 0.5) - - def test_ser_p2(self): - data = pg.GData(f"{self.dir_path:s}/twostream-f-p2.gkyl") - dg = pg.GInterpModal(data) - grid, values = dg.interpolate() - np.testing.assert_equal(len(grid[0]), 193) - np.testing.assert_equal(len(grid[1]), 97) - np.testing.assert_array_equal(values.shape, (192, 96, 1)) - np.testing.assert_approx_equal(values.mean(), 0.08337313364405809) - - def test_ser_p1_i(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - dg = pg.GInterpModal(data, poly_order=1, basis_type='ms', num_interp=3) - _, values = dg.interpolate() - np.testing.assert_array_equal(values.shape, (24, 24, 1)) - #end - - def test_ser_p2_i(self): - data = pg.GData(f"{self.dir_path:s}/twostream-f-p2.gkyl") - dg = pg.GInterpModal(data, num_interp=4) - _, values = dg.interpolate() - np.testing.assert_array_equal(values.shape, (256, 128, 1)) - #end - - def test_ten_p1(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ten-p1.gkyl") - dg = pg.GInterpModal(data, poly_order=1, basis_type="mt") - grid, values = dg.interpolate() - np.testing.assert_equal(len(grid[0]), 17) - np.testing.assert_equal(len(grid[1]), 17) - np.testing.assert_array_equal(values.shape, (16, 16, 1)) - np.testing.assert_approx_equal(values.mean(), 0.5) - - def test_ser_p1_c2p(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl", - mapc2p_name=f"{self.dir_path:s}/shock-rtheta-ser.gkyl") - dg = pg.GInterpModal(data, poly_order=1, basis_type="ms") - grid, values = dg.interpolate() - np.testing.assert_equal(len(grid[0]), 17) - np.testing.assert_equal(len(grid[1]), 17) - np.testing.assert_array_equal(values.shape, (16, 16, 1)) - np.testing.assert_approx_equal(values.mean(), 0.5) - - def test_ten_p1_c2p(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ten-p1.gkyl", - mapc2p_name=f"{self.dir_path:s}/shock-rtheta-ten.gkyl") - dg = pg.GInterpModal(data, poly_order=1, basis_type="mt") - grid, values = dg.interpolate() - np.testing.assert_equal(len(grid[0]), 17) - np.testing.assert_equal(len(grid[1]), 17) - np.testing.assert_array_equal(values.shape, (16, 16, 1)) - np.testing.assert_approx_equal(values.mean(), 0.5) diff --git a/tests/test_invariants.py b/tests/test_invariants.py new file mode 100644 index 00000000..ae564bbb --- /dev/null +++ b/tests/test_invariants.py @@ -0,0 +1,161 @@ +"""Cross-cutting laws that should hold across data shapes and backends. + +These deterministic, parametrized checks complement example-based unit tests: +they exercise algebraic, integration, representation, and state-copy contracts +over a small matrix of inputs without adding a property-testing dependency. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython + +GENERATED = Path(__file__).parent / "test_data" / "generated" +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +def _field(cells: tuple[int, ...], num_comps: int = 2) -> pg.GData: + """Build deterministic cell-average data on a nonuniform tensor grid.""" + grid = [] + for d, n in enumerate(cells): + lower, upper = -0.2 * d, 1.0 + 0.3 * d + fraction = np.linspace(0.0, 1.0, n + 1)**(d + 1) + grid.append(lower + (upper - lower) * fraction) + shape = (*cells, num_comps) + values = np.arange(np.prod(shape), dtype=float).reshape(shape) / 7.0 + 0.25 + return pg.GData().push(grid, values) + + +@pytest.mark.parametrize("cells", [(7, ), (3, 5), (2, 3, 4)]) +def test_numpy_arithmetic_obeys_affine_law_without_mutating_inputs(cells): + left = _field(cells) + right = _field(cells) + right.values[...] = np.flip(right.values, axis=0) + 0.5 + left_before = left.values.copy() + right_before = right.values.copy() + + scale = -1.75 + distributed = scale * (left + right) + expanded = scale * left + scale * right + + np.testing.assert_allclose(distributed.values, expanded.values) + np.testing.assert_array_equal(left.values, left_before) + np.testing.assert_array_equal(right.values, right_before) + assert distributed.backend == "numpy" + assert distributed.num_cells.tolist() == list(cells) + + +@pytest.mark.parametrize("ufunc", [np.negative, np.absolute, np.square, np.exp]) +def test_numpy_ufunc_result_matches_array_and_is_independent(ufunc): + source = _field((4, 3), num_comps=1) + before = source.values.copy() + + result = ufunc(source) + + np.testing.assert_allclose(result.values, ufunc(before)) + result.values.flat[0] += 100.0 + np.testing.assert_array_equal(source.values, before) + + +@pytest.mark.parametrize(("cells", "lower", "upper", "constants"), [ + ((7, ), (-2.0, ), (3.0, ), (1.25, -4.0)), + ((3, 5), (-1.0, 2.0), (2.0, 6.0), (0.5, 3.0)), + ((2, 3, 4), (0.0, -2.0, 1.0), (5.0, 1.0, 2.5), (2.0, -0.25)), +]) +def test_integral_of_constant_is_constant_times_physical_volume( + cells, lower, upper, constants): + # Uneven edges ensure the result comes from cell widths, not cell count. + grid = [] + for n, lo, hi in zip(cells, lower, upper): + fraction = np.linspace(0.0, 1.0, n + 1)**1.7 + grid.append(lo + (hi - lo) * fraction) + values = np.empty((*cells, len(constants))) + values[...] = constants + data = pg.GData().push(grid, values) + + expected = np.asarray(constants) * np.prod(np.asarray(upper) - lower) + np.testing.assert_allclose(data.integrate(), expected) + + +def test_partial_integrals_commute_and_preserve_surviving_grid(): + data = _field((3, 4, 5)) + + together = data.integrate(axis=(0, 2)) + first_then_last = data.integrate(axis=2).integrate(axis=0) + + np.testing.assert_allclose(together.values, first_then_last.values) + np.testing.assert_array_equal(together.grid[0], data.grid[1]) + assert together.num_cells.tolist() == [4] + assert data.num_cells.tolist() == [3, 4, 5] + + +@pytest.mark.parametrize("cells", [(6, ), (3, 4)]) +def test_clone_is_a_by_value_state_copy(cells): + original = _field(cells) + original.ctx["provenance"] = "original" + clone = original.clone() + + assert clone is not original + assert clone.ctx is not original.ctx + assert clone.ctx.keys() == original.ctx.keys() + np.testing.assert_array_equal(clone.num_cells, original.num_cells) + assert clone.values is not original.values + assert all(a is not b for a, b in zip(clone.grid, original.grid)) + + clone.values.flat[0] = -999.0 + clone.grid[0][0] = -999.0 + clone.ctx["provenance"] = "clone" + assert original.values.flat[0] != -999.0 + assert original.grid[0][0] != -999.0 + assert original.ctx["provenance"] == "original" + + +@needs_gkeyll +@pytest.mark.parametrize(("filename", "num_quad"), [ + ("1d_ms_p1.gkyl", 2), + ("1d_ms_p2.gkyl", 3), + ("2d_ms_p1.gkyl", 2), + ("2d_ms_p2.gkyl", 3), +]) +@pytest.mark.parametrize("value_form", ["nodal", "quad"]) +def test_native_representation_roundtrip_preserves_coefficients_and_source( + filename, num_quad, value_form): + source = pg.load(GENERATED / filename) + before = source.values.copy() + + if value_form == "nodal": + represented = source.to_nodal() + else: + represented = source.to_quad(num_quad=num_quad) + restored = represented.to_modal() + + assert represented.ctx["value_form"] == value_form + assert represented.backend == restored.backend == "gkyl" + np.testing.assert_allclose(restored.values, before, rtol=2e-13, atol=2e-13) + np.testing.assert_array_equal(source.values, before) + + +@needs_gkeyll +@pytest.mark.parametrize("value_form", ["modal", "nodal", "quad"]) +def test_native_linear_arithmetic_commutes_with_representation(value_form): + source = pg.load(GENERATED / "2d_ms_p1.gkyl") + represented = { + "modal": source, + "nodal": source.to_nodal(), + "quad": source.to_quad(), + }[value_form] + + transformed = 2.5 * represented - represented + 0.75 * represented + restored = transformed if value_form == "modal" else transformed.to_modal() + expected = 2.25 * source + + np.testing.assert_allclose(restored.values, + expected.values, + rtol=2e-13, + atol=2e-13) diff --git a/tests/test_io_h5.py b/tests/test_io_h5.py new file mode 100644 index 00000000..d6bfbb70 --- /dev/null +++ b/tests/test_io_h5.py @@ -0,0 +1,180 @@ +"""Tests for ``postgkyl.io.gkyl_h5_reader`` and ``postgkyl.io.flash_h5_reader``. + +The old test corpus has no ``.h5`` fixtures, so these build tiny files with +``tables`` directly matching each reader's expected on-disk layout (derived +from the reader source: a Gkeyll "frame" file needs a ``/StructGrid`` group +with ``vsLowerBounds``/``vsUpperBounds``/``vsNumCells`` attributes plus a +``/StructGridField`` array; a "diagnostic" file needs ``/DataStruct/timeMesh`` +and ``/DataStruct/data``; a FLASH file needs ``coordinates``/``block +size``/``node type`` plus the named field array). + +Run: PYTHONPATH=src pytest tests/test_io_h5.py -v +""" + +import os +import sys + +import numpy as np +import pytest +import tables + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl import io # noqa: E402 +from postgkyl.io.gkyl_h5_reader import GkylH5Reader # noqa: E402 +from postgkyl.io.flash_h5_reader import FlashH5Reader # noqa: E402 + + +# --------------------------------------------------------------------- gkyl h5 +def _write_gkyl_h5_frame(path, lower, upper, cells, data, time=None): + fh = tables.open_file(path, "w") + grp = fh.create_group("/", "StructGrid", "grid") + grp._v_attrs.vsLowerBounds = np.asarray(lower) + grp._v_attrs.vsUpperBounds = np.asarray(upper) + grp._v_attrs.vsNumCells = np.asarray(cells) + fh.create_array("/", "StructGridField", data) + if time is not None: + tgrp = fh.create_group("/", "timeData", "time") + tgrp._v_attrs.vsTime = time + fh.close() + + +def _write_gkyl_h5_diagnostic(path, time_mesh, data): + fh = tables.open_file(path, "w") + grp = fh.create_group("/", "DataStruct", "diag") + fh.create_array(grp, "timeMesh", time_mesh) + fh.create_array(grp, "data", data) + fh.close() + + +def test_gkyl_h5_frame_roundtrip(tmp_path): + path = str(tmp_path / "frame.h5") + data = np.arange(4 * 2 * 3, dtype=np.float64).reshape(4, 2, 3) + _write_gkyl_h5_frame(path, [0.0, -1.0], [2.0, 1.0], [4, 2], data, time=0.5) + + grid, out = io.read(path) + np.testing.assert_allclose(out, data) + assert grid[0].shape == (5, ) + assert grid[1].shape == (3, ) + np.testing.assert_allclose(grid[0], np.linspace(0.0, 2.0, 5)) + np.testing.assert_allclose(grid[1], np.linspace(-1.0, 1.0, 3)) + + +def test_gkyl_h5_frame_ctx_and_time(tmp_path): + path = str(tmp_path / "frame.h5") + data = np.ones((4, 2, 1)) + _write_gkyl_h5_frame(path, [0.0, 0.0], [1.0, 1.0], [4, 2], data, time=1.25) + + r = GkylH5Reader(path, ctx={}) + assert r.is_compatible() + r.preload() + _, out = r.load() + assert r.ctx["time"] == pytest.approx(1.25) + np.testing.assert_array_equal(r.ctx["cells"], [4, 2]) + assert r.ctx["num_comps"] == 1 + assert r.ctx["grid_type"] == "uniform" + assert out.shape == (4, 2, 1) + + +def test_gkyl_h5_diagnostic(tmp_path): + path = str(tmp_path / "diag.h5") + time_mesh = np.linspace(0.0, 1.0, 5) + data = np.arange(5 * 3, dtype=np.float64).reshape(5, 3) + _write_gkyl_h5_diagnostic(path, time_mesh, data) + + r = GkylH5Reader(path, ctx={}) + assert r.is_compatible() + assert r.is_diagnostic and not r.is_frame + r.preload() + grid, out = r.load() + np.testing.assert_allclose(out, data) + assert r.ctx["num_comps"] == 3 + assert grid[0].shape == (6, ) # uniform pseudo-grid over [time[0], time[-1]] + + +def test_gkyl_h5_is_compatible_false_for_unrelated_file(tmp_path): + path = str(tmp_path / "empty.h5") + fh = tables.open_file(path, "w") + fh.create_array("/", "SomethingElse", np.zeros(3)) + fh.close() + assert GkylH5Reader(path, ctx={}).is_compatible() is False + + +def test_gkyl_h5_is_compatible_false_for_a_non_hdf5_file(tmp_path): + path = tmp_path / "not_hdf5.h5" + path.write_bytes(b"definitely not an hdf5 file") + assert GkylH5Reader(str(path), ctx={}).is_compatible() is False + assert GkylH5Reader("/no/such/file.h5", ctx={}).is_compatible() is False + + +# ------------------------------------------------------------------- flash h5 +def _write_flash_h5(path, + *, + num_blocks=2, + nxb=4, + nyb=4, + var_name="dens", + seed=0): + """FLASH stores blocks pre-transposed on disk relative to what the reader + uses after its own ``.transpose()`` call -- see ``FlashH5Reader._read_frame``.""" + rng = np.random.default_rng(seed) + coord = np.array([[0.25, 0.25], [0.75, 0.25]][:num_blocks]) # (N, 2) + bsize = np.full((num_blocks, 2), 0.5) # (N, 2) + ntype = np.ones(num_blocks, dtype=np.int32) # (N,) all leaf blocks + bdata = rng.normal(size=(num_blocks, 1, nyb, nxb)) # -> (nxb,nyb,1,N) + + fh = tables.open_file(path, "w") + fh.create_array("/", "coordinates", coord) + fh.create_array("/", "block size", bsize) + fh.create_array("/", "node type", ntype) + fh.create_array("/", var_name, bdata) + fh.close() + + +@pytest.mark.filterwarnings( + "ignore:object name is not a valid Python identifier:tables.exceptions.NaturalNameWarning" +) +def test_flash_h5_frame_roundtrip(tmp_path): + path = str(tmp_path / "flash.h5") + _write_flash_h5(path, var_name="dens") + + r = FlashH5Reader(path, ctx={}, var_name="dens") + assert r.is_compatible() + r.preload() + grid, out = r.load() + assert out.ndim == 3 # (nx, ny, 1) + assert out.shape[-1] == 1 + assert r.ctx["grid_type"] == "uniform" + np.testing.assert_array_equal(r.ctx["cells"], out.shape[:-1]) + assert len(grid) == 2 + assert grid[0].shape == (out.shape[0] + 1, ) + assert grid[1].shape == (out.shape[1] + 1, ) + + +@pytest.mark.filterwarnings( + "ignore:object name is not a valid Python identifier:tables.exceptions.NaturalNameWarning" +) +def test_flash_h5_load_requires_var_name(tmp_path): + path = str(tmp_path / "flash.h5") + _write_flash_h5(path, var_name="dens") + r = FlashH5Reader(path, ctx={}) # var_name defaults to None + assert r.is_compatible() + r.preload() + with pytest.raises(ValueError, match="requires 'var_name'"): + r.load() + + +def test_flash_h5_is_compatible_false_without_coordinates(tmp_path): + path = str(tmp_path / "not_flash.h5") + fh = tables.open_file(path, "w") + fh.create_array("/", "SomethingElse", np.zeros(3)) + fh.close() + assert FlashH5Reader(path, ctx={}).is_compatible() is False + + +def test_flash_h5_is_compatible_false_for_a_non_hdf5_file(tmp_path): + path = tmp_path / "not_hdf5.h5" + path.write_bytes(b"definitely not an hdf5 file") + assert FlashH5Reader(str(path), ctx={}).is_compatible() is False diff --git a/tests/test_io_mapping.py b/tests/test_io_mapping.py new file mode 100644 index 00000000..f6e1aeaa --- /dev/null +++ b/tests/test_io_mapping.py @@ -0,0 +1,57 @@ +"""Tests for ``postgkyl.io.mapping``. + +Run: PYTHONPATH=src pytest tests/test_io_mapping.py -v +""" + +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +from postgkyl.io import mapping # noqa: E402 + + +def test_uniform_grid_has_cells_plus_one_edges(): + grid = mapping.uniform_grid(np.array([0.0, -1.0]), np.array([2.0, 1.0]), + np.array([4, 2])) + assert len(grid) == 2 + np.testing.assert_allclose(grid[0], [0.0, 0.5, 1.0, 1.5, 2.0]) + np.testing.assert_allclose(grid[1], [-1.0, 0.0, 1.0]) + + +def test_adjust_for_ghost_cells_no_op_when_shapes_match(): + lower = np.array([0.0]) + upper = np.array([1.0]) + cells = np.array([4]) + lo, up, c = mapping.adjust_for_ghost_cells(lower, upper, cells, (4, )) + assert c[0] == 4 + assert lo[0] == pytest.approx(0.0) + assert up[0] == pytest.approx(1.0) + + +def test_c2p_grid_splits_packed_node_axis_by_hand(): + """A hand-computed 1-D case: 3 nodes, 2 dims -> each dim gets 1 coeff.""" + # nodes[..., 0] is the x-coefficient, nodes[..., 1] the y-coefficient. + nodes = np.array([ + [0.0, 10.0], + [1.0, 11.0], + [2.0, 12.0], + ]) + blocks = mapping.c2p_grid(nodes, num_dims=2) + assert len(blocks) == 2 + np.testing.assert_array_equal(blocks[0], np.array([[0.0], [1.0], [2.0]])) + np.testing.assert_array_equal(blocks[1], np.array([[10.0], [11.0], [12.0]])) + + +def test_c2p_grid_with_multiple_coefficients_per_dim(): + """3 dims, 2 modal coefficients per dim -> 6 packed components.""" + nodes = np.arange(2 * 6, dtype=float).reshape(2, 6) + blocks = mapping.c2p_grid(nodes, num_dims=3) + assert len(blocks) == 3 + for d in range(3): + np.testing.assert_array_equal(blocks[d], nodes[:, d * 2:(d + 1) * 2]) diff --git a/tests/test_io_writer.py b/tests/test_io_writer.py new file mode 100644 index 00000000..eb520b02 --- /dev/null +++ b/tests/test_io_writer.py @@ -0,0 +1,188 @@ +"""Tests for ``postgkyl.io.writer`` -- the vtk format and series-file behavior. + +npy/txt/gkyl round trips and error paths are covered in +``tests/test_coverage_io.py``; this file focuses on what layer 04 adds: the +``vtk`` extension and its ParaView ``.series`` sidecar, plus a byte-exact +gkyl round trip through ``io.read``. + +Run: PYTHONPATH=src pytest tests/test_io_writer.py -v +""" + +import json +import os +import sys + +import numpy as np +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg # noqa: E402 +from postgkyl import io # noqa: E402 +from postgkyl.io import writer # noqa: E402 +from postgkyl.gdatastate.gdatastate import GDataState # noqa: E402 + +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F2D = os.path.join(DATA, "generated", "2d_ms_p1.gkyl") + + +def _make_state(grid, values, *, time=None, frame=None): + d = GDataState() + d.push(grid, values) + if time is not None: + d.ctx["time"] = time + if frame is not None: + d.ctx["frame"] = frame + return d + + +# --------------------------------------------------------------------- vtk +def test_vtk_writes_a_well_formed_legacy_header_1d(tmp_path): + a = pg.load(F1).interpolate().select(comp=0) + out = writer.save(a, out_name=str(tmp_path / "out1d.vtk"), extension="vtk") + assert os.path.exists(out) + with open(out, "rb") as fh: + header = fh.read(96) + assert header.startswith(b"# vtk DataFile Version") + assert b"STRUCTURED_GRID" in header + + +def test_vtk_writes_a_well_formed_legacy_header_2d(tmp_path): + b = pg.load(F2D).interpolate().select(comp=0) + out = writer.save(b, out_name=str(tmp_path / "out2d.vtk"), extension="vtk") + assert os.path.exists(out) + with open(out, "rb") as fh: + header = fh.read(96) + assert header.startswith(b"# vtk DataFile Version") + + +def test_vtk_writes_a_3d_volume(tmp_path): + grid = [ + np.linspace(0.0, 1.0, 3), + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 5) + ] + values = np.arange(2 * 3 * 4 * 1, dtype=float).reshape(2, 3, 4, 1) + d = _make_state(grid, values) + out = writer.save(d, out_name=str(tmp_path / "out3d.vtk"), extension="vtk") + assert os.path.exists(out) + with open(out, "rb") as fh: + header = fh.read(96) + assert header.startswith(b"# vtk DataFile Version") + + +def test_vtk_rejects_unsupported_dimensionality(tmp_path): + from postgkyl.io.writer import _write_vtk + grid = [np.linspace(0, 1, 2)] * 4 + values = np.ones((1, 1, 1, 1, 1)) + d = _make_state(grid, values) + with pytest.raises(ValueError, match="1-3 dimensions"): + _write_vtk(str(tmp_path / "bad.vtk"), d, 4, d.num_cells, values) + + +def test_vtk_series_file_accumulates_entries_across_two_writes(tmp_path): + grid = [np.linspace(0.0, 1.0, 4)] + values = np.array([[1.0], [2.0], [3.0]]) + + a = _make_state(grid, values, time=0.1) + out1 = writer.save(a, + out_name=str(tmp_path / "solution_0001.vtk"), + extension="vtk") + b = _make_state(grid, values, time=0.2) + out2 = writer.save(b, + out_name=str(tmp_path / "solution_0002.vtk"), + extension="vtk") + + series_path = tmp_path / "solution.vtk.series" + assert series_path.exists() + with open(series_path) as fh: + series = json.load(fh) + assert series["file-series-version"] == "1.0" + assert series["files"] == [ + { + "name": os.path.basename(out1), + "time": 0.1 + }, + { + "name": os.path.basename(out2), + "time": 0.2 + }, + ] + + +def test_vtk_series_file_updates_existing_entry_in_place(tmp_path): + """Re-writing the same frame number refreshes its time instead of + duplicating the entry.""" + grid = [np.linspace(0.0, 1.0, 4)] + values = np.array([[1.0], [2.0], [3.0]]) + + a = _make_state(grid, values, time=0.1) + writer.save(a, out_name=str(tmp_path / "solution_0001.vtk"), extension="vtk") + a2 = _make_state(grid, values, time=0.15) + writer.save(a2, out_name=str(tmp_path / "solution_0001.vtk"), extension="vtk") + + with open(tmp_path / "solution.vtk.series") as fh: + series = json.load(fh) + assert len(series["files"]) == 1 + assert series["files"][0]["time"] == pytest.approx(0.15) + + +def test_vtk_series_uses_frame_when_time_is_absent(tmp_path): + grid = [np.linspace(0.0, 1.0, 4)] + values = np.array([[1.0], [2.0], [3.0]]) + a = _make_state(grid, values, frame=3) + writer.save(a, out_name=str(tmp_path / "run_0003.vtk"), extension="vtk") + with open(tmp_path / "run.vtk.series") as fh: + series = json.load(fh) + assert series["files"][0]["time"] == pytest.approx(3.0) + + +def test_vtk_series_recovers_from_a_corrupt_sidecar(tmp_path): + grid = [np.linspace(0.0, 1.0, 4)] + values = np.array([[1.0], [2.0], [3.0]]) + (tmp_path / "bad.vtk.series").write_text("not valid json{{{") + a = _make_state(grid, values, time=0.5) + writer.save(a, out_name=str(tmp_path / "bad_0001.vtk"), extension="vtk") + with open(tmp_path / "bad.vtk.series") as fh: + series = json.load(fh) + assert len(series["files"]) == 1 + + +def test_vtk_series_recovers_from_valid_json_with_the_wrong_shape(tmp_path): + series_path = tmp_path / "bad-shape.vtk.series" + series_path.write_text("[]") + data = _make_state([np.linspace(0.0, 1.0, 4)], + np.arange(3, dtype=float)[:, None], + time=0.5) + + writer._update_vtk_series_file(data, str(tmp_path / "bad-shape_0001.vtk")) + + with open(series_path) as fh: + series = json.load(fh) + assert series["files"] == [{"name": "bad-shape_0001.vtk", "time": 0.5}] + + +# ------------------------------------------------------------------- gkyl rt +def test_gkyl_roundtrip_preserves_grid_and_values_exactly(tmp_path): + """``io.read`` is exercised both directly (grid) and through ``pg.load`` + (values, via the ``.values`` property that abstracts the gkyl/numpy + backend split -- see gdatastate/state.py) since a written already-interpolated + field still carries file_type == 1 and so is picked up again by whichever + reader is first compatible (GkylCReader when the FFI is available).""" + a = pg.load(F1).interpolate().select(comp=0) + out = writer.save(a, out_name=str(tmp_path / "rt.gkyl"), extension="gkyl") + + grid, _ = io.read(out) + for g_out, g_in in zip(a.grid, grid): + np.testing.assert_allclose(g_in, g_out) + + back = pg.load(out) + np.testing.assert_allclose(np.asarray(back.values), np.asarray(a.values)) diff --git a/tests/test_load.py b/tests/test_load.py index c09fb38d..781577b4 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -1,72 +1,205 @@ -"""Postgkyl module for testing data loading.""" -import importlib.util +"""Tests for the single-file and glob forms of ``pg.load``.""" + +from __future__ import annotations + +import importlib + import numpy as np -import os import pytest import postgkyl as pg -class TestGkyl: - """Test Gkeyll's internal binary output format.""" - dir_path = f"{os.path.dirname(__file__)}/test_data" - - def test_gkyl_type1(self): # Frame without distributed memory - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - np.testing.assert_array_equal(data.num_cells, (8, 8)) - - def test_gkyl_type1_partial(self): # Partial frame without distributed memory - data = pg.GData(f"{self.dir_path:s}/twostream-f-p2.gkyl", - z0='16', z1='8:-8', comp='0') - np.testing.assert_array_equal(data.values.shape, (1, 16, 1)) - - def test_gkyl_type1_c2p(self): # Frame with coordinate mapping - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl", - mapc2p_name=f"{self.dir_path:s}/shock-rtheta-ser.gkyl") - np.testing.assert_array_equal(data.num_cells, (8, 8)) - - def test_gkyl_type2(self): # Dynvector - data = pg.GData(f"{self.dir_path:s}/twostream-field-energy.gkyl") - np.testing.assert_array_equal(data.num_cells, (6113,)) - - def test_gkyl_type3(self): # Frame with distributed memory - data = pg.GData(f"{self.dir_path:s}/hll-euler.gkyl") - np.testing.assert_array_equal(data.num_cells, (50, 50)) - - def test_gkyl_type3_partial(self): # Partial frame with distributed memory - data = pg.GData(f"{self.dir_path:s}/hll-euler.gkyl", - z0=30, z1='30:-5', comp=0) - np.testing.assert_array_equal(data.values.shape, (1, 15, 1)) - - def test_gkyl_meta(self): # Frame with msgpack meta data included - data = pg.GData(f"{self.dir_path:s}/hll-euler.gkyl") - np.testing.assert_equal(data.ctx["frame"], 1) - - def test_gkyl_c2p_vel(self): - data = pg.GData(f"{self.dir_path:s}/bimaxwellian-elc.gkyl", - mapc2p_vel_name=f"{self.dir_path:s}/bimaxwellian-mapc2p-vel.gkyl") - dg = pg.GInterpModal(data, poly_order=1, basis_type="gkhyb") - dg.interpolate(overwrite=True) - np.testing.assert_approx_equal(data.bounds[0][1], -1.060964e07) - np.testing.assert_approx_equal(data.bounds[1][2], 1.206345e-16) -class TestAdios: - """Test Gkeyll's ADIOS2 output format.""" - dir_path = f"{os.path.dirname(__file__)}/test_data" - - adios_loader = importlib.util.find_spec('adios2') - adios_missing = adios_loader is None - - @pytest.mark.skipif(adios_missing, reason="ADIOS2 is not installed") - def test_adios_frame(self): - data = pg.GData(f"{self.dir_path:s}/twostream-f-p2_0.bp") - np.testing.assert_array_equal(data.num_cells, (64, 32)) - - @pytest.mark.skipif(adios_missing, reason="ADIOS2 is not installed") - def test_adios_frame_partial(self): - data = pg.GData(f"{self.dir_path:s}/twostream-f-p2_0.bp", z0=32, comp=0) - np.testing.assert_array_equal(data.values.shape, (1, 32, 1)) - - @pytest.mark.skipif(adios_missing, reason="ADIOS2 is not installed") - def test_adios_dynvector(self): - data = pg.GData(f"{self.dir_path:s}/twostream-field-energy.bp") - np.testing.assert_array_equal(data.num_cells, (15714,)) +class _StubData(pg.GData): + """A disk-free loaded dataset used to isolate filename dispatch.""" + + def __init__(self, file_name="", **kwargs): + super().__init__() + self._file_name = str(file_name) + self.load_kwargs = kwargs + + +@pytest.fixture +def stub_load(monkeypatch): + load_module = importlib.import_module("postgkyl.gdata.load") + monkeypatch.setattr(load_module, "GData", _StubData) + return load_module.load + + +def test_literal_filename_returns_one_dataset(stub_load): + out = stub_load("frame_0.gkyl", tag="moments") + assert isinstance(out, _StubData) + assert not isinstance(out, pg.GDataGroup) + assert out.file_name == "frame_0.gkyl" + assert out.load_kwargs["tag"] == "moments" + + +def test_pathlike_literal_remains_supported(stub_load, tmp_path): + out = stub_load(tmp_path / "frame_0.gkyl") + assert isinstance(out, _StubData) + assert out.file_name == str(tmp_path / "frame_0.gkyl") + + +def test_partial_read_options_are_lowered_once_at_load(stub_load): + out = stub_load("frame_0.gkyl", + z1="2:4", + component="1", + read_options={"custom": "yes"}) + assert out.load_kwargs["axes"] == (None, "2:4", None, None, None, None) + assert out.load_kwargs["comp"] == "1" + assert out.load_kwargs["custom"] == "yes" + + +def test_glob_returns_group_in_natural_frame_order(stub_load, tmp_path): + for frame in (10, 2, 1): + (tmp_path / f"frame_{frame}.gkyl").touch() + + out = stub_load(str(tmp_path / "frame_*.gkyl"), + label="series", + basis_type="serendipity", + poly_order=1) + + assert isinstance(out, pg.GDataGroup) + assert [d.file_name for d in out] == [ + str(tmp_path / "frame_1.gkyl"), + str(tmp_path / "frame_2.gkyl"), + str(tmp_path / "frame_10.gkyl"), + ] + assert all(d.load_kwargs["label"] == "series" for d in out) + assert all(d.load_kwargs["basis_type"] == "serendipity" for d in out) + assert all(d.load_kwargs["poly_order"] == 1 for d in out) + + +def test_single_match_glob_still_returns_group(stub_load, tmp_path): + (tmp_path / "frame_0.gkyl").touch() + out = stub_load(str(tmp_path / "frame_*.gkyl")) + assert isinstance(out, pg.GDataGroup) + assert len(out) == 1 + + +def test_unmatched_glob_has_a_targeted_error(stub_load, tmp_path): + pattern = str(tmp_path / "missing_*.gkyl") + with pytest.raises(FileNotFoundError, match="No files match pattern"): + stub_load(pattern) + + +def test_group_load_appends_and_returns_the_same_group(stub_load): + group = pg.GDataGroup() + + out = group.load("frame_0.gkyl", tag="moments").load("frame_1.gkyl") + + assert out is group + assert [data.file_name for data in group] == ["frame_0.gkyl", "frame_1.gkyl"] + assert group[0].load_kwargs["tag"] == "moments" + + +def test_group_load_appends_every_glob_match_in_natural_order( + stub_load, tmp_path): + for frame in (10, 2, 1): + (tmp_path / f"frame_{frame}.gkyl").touch() + + group = pg.GDataGroup().load(str(tmp_path / "frame_*.gkyl")) + + assert [data.file_name for data in group] == [ + str(tmp_path / "frame_1.gkyl"), + str(tmp_path / "frame_2.gkyl"), + str(tmp_path / "frame_10.gkyl"), + ] + + +def test_failed_group_load_keeps_existing_members(stub_load, tmp_path): + group = pg.GDataGroup().load("frame_0.gkyl") + + with pytest.raises(FileNotFoundError, match="No files match pattern"): + group.load(str(tmp_path / "missing_*.gkyl")) + + assert len(group) == 1 + assert group[0].file_name == "frame_0.gkyl" + + +def test_instance_load_mutates_and_returns_the_same_dataset(monkeypatch): + calls = [] + + def fake_read(file_name, ctx, **kwargs): + calls.append((file_name, kwargs)) + ctx.update(cells=np.array([2]), + basis_type="tensor", + poly_order=1, + value_form="modal", + source="reader") + return [np.linspace(0.0, 1.0, 3)], np.ones((2, 2)) + + state_module = importlib.import_module("postgkyl.gdatastate.gdatastate") + monkeypatch.setattr(state_module.io, "read", fake_read) + + data = pg.GData(tag="moments", label="ions", ctx={"seed": 7}) + out = data.load("frame_0.gkyl", + tag="loaded", + label="electrons", + basis_type="tensor", + poly_order=1, + value_form="modal", + z0=3) + + assert out is data + assert data.file_name == "frame_0.gkyl" + assert data.tag == "loaded" + assert data.label == "electrons" + assert data.ctx["seed"] == 7 + assert data.ctx["source"] == "reader" + assert data.values.shape == (2, 2) + assert calls == [("frame_0.gkyl", { + "value_form": "modal", + "basis_type": "tensor", + "poly_order": 1, + "z0": 3 + })] + + +def test_instance_reload_does_not_retain_old_file_metadata(monkeypatch): + + def fake_read(file_name, ctx, **kwargs): + ctx.update(cells=np.array([1]), + basis_type="serendipity", + poly_order=0, + value_form="nodal") + if file_name == "first.gkyl": + ctx["first_file_only"] = True + return [np.array([0.0, 1.0])], np.ones((1, 1)) + + state_module = importlib.import_module("postgkyl.gdatastate.gdatastate") + monkeypatch.setattr(state_module.io, "read", fake_read) + + data = pg.GData().load("first.gkyl") + data.load("second.gkyl") + + assert data.file_name == "second.gkyl" + assert "first_file_only" not in data.ctx + + +def test_failed_instance_load_leaves_existing_dataset_unchanged(monkeypatch): + data = pg.GData(ctx={"source": "memory"}) + data.push([np.array([0.0, 1.0])], np.array([[4.0]])) + old_grid, old_values, old_ctx = data.grid, data.values, data.ctx + + state_module = importlib.import_module("postgkyl.gdatastate.gdatastate") + + def fail_read(*args, **kwargs): + raise FileNotFoundError("missing") + + monkeypatch.setattr(state_module.io, "read", fail_read) + with pytest.raises(FileNotFoundError, match="missing"): + data.load("missing.gkyl") + + assert data.grid is old_grid + assert data.values is old_values + assert data.ctx is old_ctx + assert data.file_name == "" + + +def test_instance_load_rejects_empty_names_and_globs(): + data = pg.GData() + with pytest.raises(ValueError, match="non-empty filename"): + data.load("") + with pytest.raises(ValueError, match=r"pg\.load\(pattern\)"): + data.load("frame_*.gkyl") diff --git a/tests/test_multiblock.py b/tests/test_multiblock.py new file mode 100644 index 00000000..48ed75f7 --- /dev/null +++ b/tests/test_multiblock.py @@ -0,0 +1,327 @@ +"""Native multiblock support: identity, partition, and one-figure terminals. + +Gkeyll writes a decomposed-domain run as one file per block, +``'_b-_.gkyl'``. Those files are *one field*, so +postgkyl must (a) recognize the block index without being told, (b) keep +verbs acting blockwise, and (c) have terminal verbs act on the field as a +whole -- one figure, one color scale, one colorbar. + +The fixtures are the ``mb_sim_b{0,1,2}-elc_M0_{0,1}`` family written by +``generate_test_data.generate_all``: three blocks tiling the x axis into +abutting disjoint domains, two frames each. +""" + +from __future__ import annotations + +import glob +import os + +import matplotlib.figure +import matplotlib.pyplot as plt +import numpy as np +import pytest +from click.testing import CliRunner + +import postgkyl as pg +from postgkyl.cli.app import cli +from postgkyl.io import parse_output_name + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +GEN = os.path.join(ROOT, "tests", "test_data", "generated") +MB_GLOB = os.path.join(GEN, "mb_sim_b*-elc_M0_0.gkyl") +MB_GLOB_ALL_FRAMES = os.path.join(GEN, "mb_sim_b*-elc_M0_*.gkyl") + + +def _run(args): + return CliRunner().invoke(cli, args) + + +def _ok(args): + result = _run(args) + assert result.exit_code == 0, result.output + return result + + +def _blocks(frame: int = 0): + return [ + pg.load(fn) for fn in sorted( + glob.glob(os.path.join(GEN, f"mb_sim_b*-elc_M0_{frame}.gkyl"))) + ] + + +# ============================================================ the parser +class TestParseOutputName: + + @pytest.mark.parametrize( + "name, sim, block, quantity, frame", + [ + # The real file the convention was confirmed against. + ("rt_gk_multib_sheath_1x2v_p1_b2-geo_int_B3.gkyl", + "rt_gk_multib_sheath_1x2v_p1", 2, "geo_int_B3", None), + ("sim_b10-elc_M0_7.gkyl", "sim", 10, "elc_M0", 7), + ("gk_lorentzian_mirror-elc_M0_1.gkyl", "gk_lorentzian_mirror", None, + "elc_M0", 1), + ("sim-dt.gkyl", "sim", None, "dt", None), + ]) + def test_parses_the_convention(self, name, sim, block, quantity, frame): + parsed = parse_output_name(name) + assert (parsed.sim, parsed.block, parsed.quantity, + parsed.frame) == (sim, block, quantity, frame) + + def test_frame_requires_all_digits(self): + # 'geo_int_B3' must not be read as quantity 'geo_int_B' at frame 3 -- + # the trailing run has to be digits only. + assert parse_output_name("sim_b0-geo_int_B3.gkyl").frame is None + + def test_block_requires_digits_so_a_sim_named__b_is_safe(self): + # A simulation legitimately named '..._beta' must not be read as block + # 'eta': _b needs digits. + parsed = parse_output_name("gk_beta-elc_M0_0.gkyl") + assert parsed.block is None + assert parsed.sim == "gk_beta" + + def test_prefix_is_per_block_so_geometry_resolves_per_block(self): + assert parse_output_name("d/sim_b2-elc_M0_3.gkyl").prefix == "d/sim_b2" + assert parse_output_name("d/sim-elc_M0_3.gkyl").prefix == "d/sim" + + def test_restart_suffix_is_stripped(self): + parsed = parse_output_name("sim-elc_5_restart.gkyl") + assert (parsed.quantity, parsed.frame, parsed.restart) == ("elc", 5, True) + + def test_field_key_excludes_block_and_directory(self): + parsed = parse_output_name("dir/sim_b2-elc_M0_3.gkyl") + assert parsed.field_key == ("sim", "elc_M0", 3) + + def test_empty_path_has_no_identity(self): + assert parse_output_name("") is None + assert parse_output_name(None) is None + + +# ====================================================== identity in ctx +class TestBlockIdentityIsStamped: + + def test_load_stamps_sim_block_quantity_frame(self): + data = pg.load(os.path.join(GEN, "mb_sim_b1-elc_M0_0.gkyl")) + assert data.ctx["sim"] == "mb_sim" + assert data.ctx["block"] == 1 + assert data.ctx["quantity"] == "elc_M0" + assert data.ctx["frame"] == 0 + + def test_single_block_data_has_block_none(self): + data = pg.load(os.path.join(GEN, "distf_p2_0.gkyl")) + assert data.ctx["block"] is None + + def test_identity_survives_verbs(self): + # clone() copies ctx, so a family is still recognizable downstream -- + # this is what lets 'interp ... plot' still draw the blocks together. + out = pg.load(os.path.join(GEN, "mb_sim_b2-elc_M0_1.gkyl")).interpolate() + assert out.ctx["block"] == 2 + assert out.ctx["frame"] == 1 + + def test_header_frame_wins_over_the_file_name(self): + # The reader stamps frame from the file's msgpack metadata; the parsed + # name only fills a gap, never overrides. + data = pg.load(os.path.join(GEN, "mb_sim_b0-elc_M0_1.gkyl")) + assert data.ctx["frame"] == 1 + + def test_identity_is_not_written_into_saved_files(self, tmp_path): + # The identity comes from the *path*, so it must never be stored in the + # file: saving block 1's data under another name and reloading it would + # otherwise find a stale block index in the header -- which, because + # header metadata wins over the parsed name, would silently stick. + data = pg.load(os.path.join(GEN, "mb_sim_b1-elc_M0_0.gkyl")).interpolate() + out = pg.save(data, out_name=str(tmp_path / "plain-thing_0.gkyl")) + reloaded = pg.load(out) + assert reloaded.ctx["block"] is None + assert reloaded.ctx["sim"] == "plain" + assert reloaded.ctx["quantity"] == "thing" + + def test_info_reports_the_block(self): + out = pg.load(os.path.join(GEN, "mb_sim_b1-elc_M0_0.gkyl")).info() + assert "Block: 1" in out + # The identity keys must not also fall through to info's generic ctx dump. + assert "├─ block:" not in out + assert "├─ sim:" not in out + + +# ========================================================== the partition +class TestGroupBlocks: + + def test_one_family_per_field(self): + families = pg.group_blocks(_blocks()) + assert len(families) == 1 + assert [d.ctx["block"] for d in families[0]] == [0, 1, 2] + + def test_frames_are_separate_families(self): + both = _blocks(0) + _blocks(1) + families = pg.group_blocks(both) + assert len(families) == 2 + assert {d.ctx["frame"] for d in families[0]} == {0} + assert {d.ctx["frame"] for d in families[1]} == {1} + + def test_family_is_sorted_by_block_index(self): + shuffled = list(reversed(_blocks())) + assert [d.ctx["block"] for d in pg.group_blocks(shuffled)[0]] == [0, 1, 2] + + def test_single_block_data_is_all_singletons(self): + # The property that keeps every pre-existing pipeline unchanged. + frames = [pg.load(os.path.join(GEN, f"distf_p2_{i}.gkyl")) for i in (0, 1)] + assert pg.group_blocks(frames) == [[frames[0]], [frames[1]]] + + def test_differently_tagged_results_do_not_merge(self): + blocks = _blocks() + tagged = [d.interpolate(tag="rz") for d in blocks] + families = pg.group_blocks(blocks + tagged) + assert len(families) == 2 + assert {d.tag for d in families[0]} == {"default"} + assert {d.tag for d in families[1]} == {"rz"} + + +# ================================================= terminals: one figure +class TestOneFigurePerField: + + def test_group_plot_is_one_figure(self): + group = pg.GDataGroup([d.interpolate() for d in _blocks()]) + fig = group.plot(no_show=True) + assert isinstance(fig, matplotlib.figure.Figure) + + def test_operations_plot_takes_many_datasets(self): + fig = pg.plot(*[d.interpolate() for d in _blocks()], no_show=True) + assert isinstance(fig, matplotlib.figure.Figure) + + def test_blocks_share_one_color_scale_and_one_colorbar(self): + # Each block's values are offset by its block index, so per-dataset + # normalization would give three different scales for one field. + blocks = [d.interpolate() for d in _blocks()] + fig = pg.plot(*blocks, no_show=True) + ax = fig.axes[0] + meshes = ax.collections + assert len(meshes) == 3 + clims = {m.get_clim() for m in meshes} + assert len(clims) == 1, f"blocks drew on different color scales: {clims}" + + expected = (min(float(np.nanmin(d.values)) for d in blocks), + max(float(np.nanmax(d.values)) for d in blocks)) + assert clims.pop() == pytest.approx(expected) + + # One colorbar for the panel, not one per block: the pcolormesh panel + # plus a single appended colorbar axes. + assert len(fig.axes) == 2 + + def test_explicit_zlim_still_wins(self): + blocks = [d.interpolate() for d in _blocks()] + fig = pg.plot(*blocks, zmin=-1.0, zmax=1.0, no_show=True) + for mesh in fig.axes[0].collections: + assert mesh.get_clim() == pytest.approx((-1.0, 1.0)) + + +# ================================== per-block geometry (gk_rz/gk_fluxsurf) +class TestPerBlockGeometry: + + def test_geometry_prefix_is_per_block(self): + from postgkyl.diagnostics.gk import rz + + assert rz.geometry_prefix("d/sim_b2-elc_M0_3.gkyl") == "d/sim_b2" + assert rz.geometry_prefix("d/sim-elc_M0_3.gkyl") == "d/sim" + assert rz.geometry_prefix("") is None + + def test_explicit_geometry_path_substitutes_the_block_index(self): + from postgkyl.diagnostics.gk.rz import per_block_path + + assert per_block_path("geo_b*.gkyl", 3) == "geo_b3.gkyl" + assert per_block_path("geo.gkyl", 3) == "geo.gkyl" # no '*' -> as given + assert per_block_path("geo_b*.gkyl", None) == "geo_b*.gkyl" # single block + assert per_block_path(None, 3) is None + + def test_each_block_resolves_its_own_geometry(self, monkeypatch): + # The bug this replaces: geometry was resolved once, from the first + # dataset, and that one projection was applied to every block -- drawing + # every block at block 0's position. + from postgkyl.diagnostics.gk import rz + + seen = [] + monkeypatch.setattr( + rz, "resolve_geometry", + lambda file_name, **kw: seen.append(file_name) or file_name) + monkeypatch.setattr(rz, "resolve_rz_projection", lambda first, geo, **kw: + ("projection", geo)) + + blocks = _blocks(0) + _blocks(1) # 3 blocks x 2 frames + projections = rz.rz_projections(blocks) + + # One geometry read per block, not per dataset and not just one overall. + assert len(seen) == 3 + assert set(projections) == { + os.path.join(GEN, f"mb_sim_b{b}") + for b in (0, 1, 2) + } + for data in blocks: + assert rz.projection_for( + projections, data) is projections[rz.geometry_prefix(data.file_name)] + + def test_interpolated_grid_values_is_idempotent(self): + # 'pgkyl ... interp gk_rz' must not interpolate twice: the second pass + # would run the DG evaluation matrix over values that are already point + # values, silently producing garbage instead of raising. + from postgkyl.diagnostics.gk import utils + + raw = pg.load(os.path.join(GEN, "mb_sim_b0-elc_M0_0.gkyl")) + once = utils.interpolated_grid_values(raw) + twice = utils.interpolated_grid_values(raw.interpolate()) + assert np.allclose(once[2], twice[2]) + + +# ================================================================== CLI +class TestMultiblockCli: + + @staticmethod + def _plot_calls(monkeypatch): + """Record figures made by the one canonical plot call.""" + calls = [] + real = plt.figure + + def spy(*args, **kwargs): + figure = real(*args, **kwargs) + calls.append(figure) + return figure + + monkeypatch.setattr(plt, "figure", spy) + return calls + + def test_plot_draws_all_blocks_on_one_figure(self, monkeypatch): + calls = self._plot_calls(monkeypatch) + _ok([MB_GLOB, "interp", "plot", "--no_show"]) + assert len(calls) == 1 + assert len(calls[0].axes[0].collections) == 3 + + def test_two_frames_give_two_figures(self, monkeypatch): + calls = self._plot_calls(monkeypatch) + _ok([MB_GLOB_ALL_FRAMES, "interp", "plot", "--no_show"]) + assert len(calls) == 2 + assert all(len(figure.axes[0].collections) == 3 for figure in calls) + + def test_single_block_data_still_gets_a_figure_per_dataset(self, monkeypatch): + calls = self._plot_calls(monkeypatch) + _ok([os.path.join(GEN, "distf_p2_*.gkyl"), "interp", "plot", "--no_show"]) + assert len(calls) == 2 + + def test_multiblock_flag_forces_everything_onto_one_figure(self, monkeypatch): + calls = self._plot_calls(monkeypatch) + _ok([MB_GLOB_ALL_FRAMES, "interp", "plot", "--multiblock", "--no_show"]) + assert len(calls) == 1 + assert len(calls[0].axes[0].collections) == 6 + + def test_load_orders_blocks_naturally(self): + # A lexicographic sort would put _b10 before _b2; the working set must + # be in block order. + result = _ok([MB_GLOB, "info"]) + blocks = [ + int(line.split(":")[1].split("(")[0]) + for line in result.output.splitlines() if line.startswith("├─ Block:") + ] + assert blocks == [0, 1, 2] + + def test_explicit_render_options_save_one_png_for_the_field(self, tmp_path): + out = tmp_path / "mb" + _ok([MB_GLOB, "interp", "plot", "--no_show", "--saveas", str(out)]) + assert sorted(p.name for p in tmp_path.glob("*.png")) == ["mb.png"] diff --git a/tests/test_numerics_calculus.py b/tests/test_numerics_calculus.py new file mode 100644 index 00000000..e6d0946f --- /dev/null +++ b/tests/test_numerics_calculus.py @@ -0,0 +1,129 @@ +"""Tests for postgkyl.numerics.calculus -- integrate over a nodal grid.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.numerics import calculus + + +class TestIntegrate1D: + + def test_uniform_ones_integrates_to_domain_length(self): + grid = [np.linspace(0.0, 1.0, 6)] # 5 cells, dx=0.2 + _, out = calculus.integrate(grid, np.ones((5, 1)), axis=0) + np.testing.assert_allclose(out.flat[0], 1.0, rtol=1e-12) + + def test_linear_function_exact_integral(self): + # integral of x from 0 to 1 = 0.5 (analytic, hand-computed) + N = 100 + grid = [np.linspace(0.0, 1.0, N + 1)] + x_cc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = x_cc[:, np.newaxis] + _, out = calculus.integrate(grid, values, axis=0) + np.testing.assert_allclose(out.flat[0], 0.5, rtol=1e-3) + + def test_quadratic_function_exact_integral(self): + # integral of x^2 from 0 to 1 = 1/3 (analytic, hand-computed) + N = 4000 + grid = [np.linspace(0.0, 1.0, N + 1)] + x_cc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = (x_cc**2)[:, np.newaxis] + _, out = calculus.integrate(grid, values, axis=0) + np.testing.assert_allclose(out.flat[0], 1.0 / 3.0, rtol=1e-3) + + def test_integer_axis(self): + grid = [np.linspace(0.0, 2.0, 5)] # 4 cells, dx=0.5 + _, out = calculus.integrate(grid, np.ones((4, 1)), axis=0) + np.testing.assert_allclose(out.flat[0], 2.0, rtol=1e-12) + + def test_string_integer_axis(self): + grid = [np.linspace(0.0, 1.0, 6)] + _, out = calculus.integrate(grid, np.ones((5, 1)), axis="0") + np.testing.assert_allclose(out.flat[0], 1.0, rtol=1e-12) + + def test_tuple_axis(self): + grid = [np.linspace(0.0, 1.0, 6)] + _, out = calculus.integrate(grid, np.ones((5, 1)), axis=(0, )) + np.testing.assert_allclose(out.flat[0], 1.0, rtol=1e-12) + + def test_none_axis_integrates_all(self): + grid = [np.linspace(0.0, 1.0, 6)] + _, out = calculus.integrate(grid, np.ones((5, 1)), axis=None) + np.testing.assert_allclose(out.flat[0], 1.0, rtol=1e-12) + + def test_colon_slice_axis_string(self): + """src_bak's colon-slice branch passed raw strings to ``range()``, + which raises TypeError immediately -- a latent bug never exercised by + any caller. Fixed here (cast to int) and proven by this test.""" + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] + _, out = calculus.integrate(grid, np.ones((5, 4, 1)), axis="0:2") + np.testing.assert_allclose(out.flat[0], 2.0, rtol=1e-12) + + def test_does_not_mutate_input_values(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + calculus.integrate(grid, values, axis=0) + np.testing.assert_allclose(values, np.ones((5, 1))) + + def test_wrong_axis_type_raises(self): + grid = [np.linspace(0.0, 1.0, 6)] + with pytest.raises(TypeError): + calculus.integrate(grid, np.ones((5, 1)), axis=3.14) + + def test_output_shape_preserved_with_expand_dims(self): + grid = [np.linspace(0.0, 1.0, 6)] + _, out = calculus.integrate(grid, np.ones((5, 2)), axis=0) + assert out.shape == (1, 2) + + def test_multiple_components(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.column_stack([np.ones(5), 2.0 * np.ones(5)]) + _, out = calculus.integrate(grid, values, axis=0) + np.testing.assert_allclose(out[0, 0], 1.0, rtol=1e-12) + np.testing.assert_allclose(out[0, 1], 2.0, rtol=1e-12) + + +class TestIntegrate2D: + + def test_ones_integrates_to_area(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] # 5x4 cells + _, out = calculus.integrate(grid, np.ones((5, 4, 1)), axis=None) + np.testing.assert_allclose(out.flat[0], 2.0, rtol=1e-12) + + def test_integrate_axis0_only(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 1.0, 4)] # 5x3 + _, out = calculus.integrate(grid, np.ones((5, 3, 1)), axis=0) + assert out.shape == (1, 3, 1) + np.testing.assert_allclose(out[:, :, 0], 1.0, rtol=1e-12) + + def test_integrate_axis1_only(self): + grid = [np.linspace(0.0, 1.0, 4), np.linspace(0.0, 2.0, 5)] # 3x4 + _, out = calculus.integrate(grid, np.ones((3, 4, 1)), axis=1) + assert out.shape == (3, 1, 1) + np.testing.assert_allclose(out[:, :, 0], 2.0, rtol=1e-12) + + def test_comma_separated_string_axes(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] + _, out = calculus.integrate(grid, np.ones((5, 4, 1)), axis="0,1") + np.testing.assert_allclose(out.flat[0], 2.0, rtol=1e-12) + + def test_nonuniform_grid(self): + x = np.array([0.0, 0.1, 0.4, 1.0]) + _, out = calculus.integrate([x], np.ones((3, 1)), axis=0) + np.testing.assert_allclose(out.flat[0], 1.0, rtol=1e-12) + + +class TestIntegrateCellCentered: + + def test_cell_centered_grid(self): + # len(coord) == values.shape[d] -> a last element is appended to dz + x_cc = np.linspace(0.1, 0.9, 5) # 5 cell centers, dx=0.2 + _, out = calculus.integrate([x_cc], np.ones((5, 1)), axis=0) + np.testing.assert_allclose(out.flat[0], 1.0, rtol=1e-12) + + def test_single_cell_axis_uses_mean(self): + grid = [np.array([0.5]), np.linspace(0.0, 1.0, 4)] + _, out = calculus.integrate(grid, np.ones((1, 3, 1)), axis=0) + assert out.shape[0] == 1 diff --git a/tests/test_numerics_downsample.py b/tests/test_numerics_downsample.py new file mode 100644 index 00000000..8193f671 --- /dev/null +++ b/tests/test_numerics_downsample.py @@ -0,0 +1,79 @@ +"""Tests for postgkyl.numerics.downsample.""" + +from __future__ import annotations + +import numpy as np + +from postgkyl.numerics.downsample import downsample + + +class TestDownsample: + + def test_no_arrays_returns_empty_tuple(self): + assert downsample() == () + + def test_zero_max_points_returns_unchanged(self): + x = np.linspace(0, 1, 100) + out, = downsample(x, maximum_points_per_axis=0) + np.testing.assert_array_equal(out, x) + + def test_negative_max_points_returns_unchanged(self): + x = np.linspace(0, 1, 100) + out, = downsample(x, maximum_points_per_axis=-5) + np.testing.assert_array_equal(out, x) + + def test_none_max_points_returns_unchanged(self): + x = np.linspace(0, 1, 100) + out, = downsample(x, maximum_points_per_axis=None) + np.testing.assert_array_equal(out, x) + + def test_scalar_array_returns_unchanged(self): + x = np.array(5.0) + out, = downsample(x, maximum_points_per_axis=10) + np.testing.assert_array_equal(out, x) + + def test_mismatched_shapes_returns_unchanged(self): + x = np.linspace(0, 1, 100) + y = np.linspace(0, 1, 50) + out_x, out_y = downsample(x, y, maximum_points_per_axis=10) + assert out_x.shape == (100, ) + assert out_y.shape == (50, ) + + def test_already_within_limit_returns_unchanged(self): + x = np.linspace(0, 1, 5) + out, = downsample(x, maximum_points_per_axis=20) + np.testing.assert_array_equal(out, x) + + def test_1d_downsampling_caps_axis_length(self): + x = np.linspace(0, 10, 100) + out, = downsample(x, maximum_points_per_axis=20) + assert out.shape[0] <= 21 + + def test_1d_downsampling_keeps_endpoints(self): + x = np.linspace(0, 10, 100) + out, = downsample(x, maximum_points_per_axis=20) + assert out[0] == x[0] + assert out[-1] == x[-1] + + def test_downsampling_does_not_duplicate_an_aligned_endpoint(self): + x = np.arange(5) + out, = downsample(x, maximum_points_per_axis=3) + np.testing.assert_array_equal(out, [0, 2, 4]) + + def test_multiple_arrays_downsampled_consistently(self): + x = np.linspace(0, 10, 100) + y = np.sin(x) + x_ds, y_ds = downsample(x, y, maximum_points_per_axis=10) + assert x_ds.shape == y_ds.shape + np.testing.assert_allclose(y_ds, np.sin(x_ds)) + + def test_2d_downsampling(self): + value = np.random.default_rng(0).random((100, 100)) + out, = downsample(value, maximum_points_per_axis=10) + assert out.shape[0] <= 11 + assert out.shape[1] <= 11 + + def test_3d_downsampling(self): + value = np.random.default_rng(0).random((30, 30, 30)) + out, = downsample(value, maximum_points_per_axis=10) + assert all(s <= 11 for s in out.shape) diff --git a/tests/test_numerics_ev_ops.py b/tests/test_numerics_ev_ops.py new file mode 100644 index 00000000..87b6c1c4 --- /dev/null +++ b/tests/test_numerics_ev_ops.py @@ -0,0 +1,466 @@ +"""Tests for postgkyl.numerics.ev_ops -- the RPN operator registry. + +Every operator is ``f(in_grid, in_values) -> ([out_grid], [out_values])`` +over plain lists / NumPy arrays. ``cmds`` maps each RPN token to its +arity and function. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.numerics import ev_ops + + +def _arr(*vals): + return np.array(vals, dtype=float) + + +class TestCmdsTable: + + def test_expected_keys_present(self): + expected = { + "+", + "-", + "*", + "/", + "dot", + "sqrt", + "sin", + "cos", + "tan", + "abs", + "avg", + "log", + "log10", + "max", + "min", + "max2", + "min2", + "mean", + "len", + "pow", + "sq", + "exp", + "grad", + "grad2", + "int", + "div", + "curl", + "scale_comp", + "scale_zi_axis", + } + assert set(ev_ops.cmds) == expected + + def test_arities_are_ints(self): + for tok, spec in ev_ops.cmds.items(): + assert isinstance(spec["num_in"], int) + assert isinstance(spec["num_out"], int) + assert callable(spec["func"]) + + +class TestGetGrid: + + def test_both_none(self): + assert ev_ops._get_grid(None, None) is None + + def test_first_none(self): + g = [np.array([0.0, 1.0])] + assert ev_ops._get_grid(None, g) is g + + def test_second_none(self): + g = [np.array([0.0, 1.0])] + assert ev_ops._get_grid(g, None) is g + + def test_prefers_longer_grid(self): + g1 = [np.array([0.0, 1.0])] + g2 = [np.array([0.0, 1.0]), np.array([0.0, 1.0])] + assert ev_ops._get_grid(g1, g2) is g2 + assert ev_ops._get_grid(g2, g1) is g2 + + +class TestArithmetic: + + def test_add(self): + out_grid, out_vals = ev_ops.add([None, None], [_arr(1.0), _arr(2.0)]) + np.testing.assert_allclose(out_vals[0], 3.0) + + def test_subtract_is_stack_order(self): + # RPN: a b - computes b - a (in_values[1] - in_values[0]) + _, out_vals = ev_ops.subtract([None, None], [_arr(1.0), _arr(5.0)]) + np.testing.assert_allclose(out_vals[0], 4.0) + + def test_mult_same_shape(self): + _, out_vals = ev_ops.mult([None, None], [_arr(2.0), _arr(3.0)]) + np.testing.assert_allclose(out_vals[0], 6.0) + + def test_mult_broadcast_leading_axis(self): + """Cross-basis (conf x phase) multiply: the conf-space field's leading + axis matches the phase-space field's leading axis, so multiply via + transpose-multiply-transpose instead of NumPy's trailing-axis rule.""" + conf = np.ones((3, 1)) # 3 conf cells, 1 comp + phase = np.arange(12.0).reshape(3, 4) # 3 conf cells x 4 vel cells + _, out_vals = ev_ops.mult([None, None], [conf, phase]) + expected = (phase.transpose() * conf.transpose()).transpose() + np.testing.assert_allclose(out_vals[0], expected) + + def test_divide_same_shape(self): + _, out_vals = ev_ops.divide([None, None], [_arr(2.0), _arr(10.0)]) + np.testing.assert_allclose(out_vals[0], 5.0) + + def test_divide_broadcast_leading_axis(self): + conf = np.full((3, 1), 2.0) + phase = np.arange(1.0, 13.0).reshape(3, 4) + _, out_vals = ev_ops.divide([None, None], [conf, phase]) + expected = (phase.transpose() / conf.transpose()).transpose() + np.testing.assert_allclose(out_vals[0], expected) + + def test_dot(self): + g = [np.array([0.0, 1.0])] + a = np.array([[1.0, 2.0, 3.0]]) + b = np.array([[4.0, 5.0, 6.0]]) + out_grid, out_vals = ev_ops.dot([g, g], [a, b]) + np.testing.assert_allclose(out_vals[0], [[32.0]]) + + +class TestUnaryMath: + + def test_sqrt(self): + g = [np.array([0.0, 1.0])] + _, out_vals = ev_ops.sqrt([g], [_arr(4.0)]) + np.testing.assert_allclose(out_vals[0], 2.0) + + def test_sin_cos_tan(self): + g = [np.array([0.0, 1.0])] + x = _arr(0.0) + np.testing.assert_allclose(ev_ops.psin([g], [x])[1][0], 0.0) + np.testing.assert_allclose(ev_ops.pcos([g], [x])[1][0], 1.0) + np.testing.assert_allclose(ev_ops.ptan([g], [x])[1][0], 0.0) + + def test_absolute(self): + g = [np.array([0.0, 1.0])] + _, out_vals = ev_ops.absolute([g], [_arr(-3.0)]) + np.testing.assert_allclose(out_vals[0], 3.0) + + def test_log_and_log10(self): + g = [np.array([0.0, 1.0])] + np.testing.assert_allclose(ev_ops.log([g], [_arr(np.e)])[1][0], 1.0) + np.testing.assert_allclose(ev_ops.log10([g], [_arr(100.0)])[1][0], 2.0) + + def test_sq(self): + g = [np.array([0.0, 1.0])] + _, out_vals = ev_ops.sq([g], [_arr(3.0)]) + np.testing.assert_allclose(out_vals[0], 9.0) + + def test_exp(self): + g = [np.array([0.0, 1.0])] + _, out_vals = ev_ops.exp([g], [_arr(0.0)]) + np.testing.assert_allclose(out_vals[0], 1.0) + + +class TestReductions: + + def test_minimum(self): + _, out_vals = ev_ops.minimum([None], [np.array([3.0, 1.0, 2.0])]) + np.testing.assert_allclose(out_vals[0], [1.0]) + + def test_minimum_ignores_nan(self): + _, out_vals = ev_ops.minimum([None], [np.array([np.nan, 1.0, 2.0])]) + np.testing.assert_allclose(out_vals[0], [1.0]) + + def test_maximum(self): + _, out_vals = ev_ops.maximum([None], [np.array([3.0, 1.0, 2.0])]) + np.testing.assert_allclose(out_vals[0], [3.0]) + + def test_mean(self): + _, out_vals = ev_ops.mean([None], [np.array([1.0, 2.0, 3.0])]) + np.testing.assert_allclose(out_vals[0], [2.0]) + + def test_minimum2(self): + _, out_vals = ev_ops.minimum2( + [None, None], [_arr(1.0, 5.0), _arr(3.0, 2.0)]) + np.testing.assert_allclose(out_vals[0], [1.0, 2.0]) + + def test_maximum2(self): + _, out_vals = ev_ops.maximum2( + [None, None], [_arr(1.0, 5.0), _arr(3.0, 2.0)]) + np.testing.assert_allclose(out_vals[0], [3.0, 5.0]) + + +class TestPower: + + def test_power_is_stack_order(self): + # RPN: a b pow computes b ** a (in_values[1] ** in_values[0]) + _, out_vals = ev_ops.power([None, _arr(0.0)], [_arr(2.0), _arr(3.0)]) + np.testing.assert_allclose(out_vals[0], 9.0) + + +class TestLength: + + def test_nodal_grid_length(self): + grid = [np.linspace(0.0, 4.0, 5)] # nodal, 4 cells + values = np.ones((4, 1)) + _, out_vals = ev_ops.length([None, grid], [0.0, values]) + np.testing.assert_allclose(out_vals[0], 4.0) + + def test_cell_centered_grid_length_adds_one_more_dz(self): + """When ``len(coord) == values.shape[axis]`` (already-cell-centered + grid), one extra spacing is added, matching the ``calculus.integrate`` + convention.""" + grid = [np.linspace(0.0, 3.0, 4)] # 4 cell centers, dx=1 + values = np.ones((4, 1)) + _, out_vals = ev_ops.length([None, grid], [0.0, values]) + np.testing.assert_allclose(out_vals[0], 4.0) + + +class TestGrad: + + def test_grad_1d_matches_analytic_slope(self): + grid = [np.linspace(0.0, 1.0, 11)] # 10 cells + zc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = (2.0 * zc)[:, np.newaxis] # f(x) = 2x -> df/dx = 2 + _, out_vals = ev_ops.grad([grid], [values]) + np.testing.assert_allclose(out_vals[0][:, 0], 2.0, rtol=1e-8) + + def test_grad2_colon_range(self): + grid = [np.linspace(0.0, 1.0, 11), np.linspace(0.0, 1.0, 11)] + values = np.ones((10, 10, 1)) + _, out_vals = ev_ops.grad2([None, grid], ["0:2", values]) + assert out_vals[0].shape[-1] == 2 + + def test_grad2_comma_list(self): + grid = [np.linspace(0.0, 1.0, 11), np.linspace(0.0, 1.0, 11)] + values = np.ones((10, 10, 1)) + _, out_vals = ev_ops.grad2([None, grid], ["0,1", values]) + assert out_vals[0].shape[-1] == 2 + + def test_grad2_single_axis(self): + grid = [np.linspace(0.0, 1.0, 11)] + zc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = (3.0 * zc)[:, np.newaxis] + _, out_vals = ev_ops.grad2([None, grid], [0, values]) + np.testing.assert_allclose(out_vals[0][:, 0], 3.0, rtol=1e-8) + + +class TestIntegrateAndAverage: + + def test_integrate_matches_calculus_integrate(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + _, out_vals = ev_ops.integrate([None, grid], [np.array(0.0), values]) + np.testing.assert_allclose(out_vals[0].flat[0], 1.0, rtol=1e-12) + + def test_integrate_axis_all_string(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] + values = np.ones((5, 4, 1)) + _, out_vals = ev_ops.integrate([None, grid], ["all", values]) + np.testing.assert_allclose(out_vals[0].flat[0], 2.0, rtol=1e-12) + + def test_integrate_colon_slice_axis(self): + """src_bak's colon-slice branch passed raw strings to ``range()``, + a TypeError-raising latent bug never exercised by any caller. Fixed + here (cast to int) and proven by this test.""" + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] + values = np.ones((5, 4, 1)) + _, out_vals = ev_ops.integrate([None, grid], ["0:2", values]) + np.testing.assert_allclose(out_vals[0].flat[0], 2.0, rtol=1e-12) + + def test_integrate_ndarray_axis(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + _, out_vals = ev_ops.integrate([None, grid], [np.array(0.0), values]) + np.testing.assert_allclose(out_vals[0].flat[0], 1.0, rtol=1e-12) + + def test_integrate_bad_axis_type_raises(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + with pytest.raises(TypeError): + ev_ops.integrate([None, grid], [3 + 4j, values]) + + def test_average_divides_by_length(self): + grid = [np.linspace(0.0, 2.0, 6)] # length 2, 5 cells + values = 3.0 * np.ones((5, 1)) + _, out_vals = ev_ops.average([None, grid], [np.array(0.0), values]) + np.testing.assert_allclose(out_vals[0].flat[0], 3.0, rtol=1e-10) + + def test_integrate_float_axis(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + _, out_vals = ev_ops.integrate([None, grid], [0.0, values]) + np.testing.assert_allclose(out_vals[0].flat[0], 1.0, rtol=1e-12) + + def test_integrate_tuple_axis(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] + values = np.ones((5, 4, 1)) + _, out_vals = ev_ops.integrate([None, grid], [(0, 1), values]) + np.testing.assert_allclose(out_vals[0].flat[0], 2.0, rtol=1e-12) + + def test_integrate_comma_string_axis(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 2.0, 5)] + values = np.ones((5, 4, 1)) + _, out_vals = ev_ops.integrate([None, grid], ["0,1", values]) + np.testing.assert_allclose(out_vals[0].flat[0], 2.0, rtol=1e-12) + + def test_integrate_single_int_string_axis(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + _, out_vals = ev_ops.integrate([None, grid], ["0", values]) + np.testing.assert_allclose(out_vals[0].flat[0], 1.0, rtol=1e-12) + + def test_integrate_cell_centered_grid_appends_last_spacing(self): + """When ``len(coord) == values.shape[d]`` (an already-cell-centered + grid), one extra spacing is appended to ``dz`` -- matching + ``calculus.integrate``'s convention.""" + x_cc = np.linspace(0.1, 0.9, 5) # 5 cell centers, dx=0.2 + _, out_vals = ev_ops.integrate( + [None, [x_cc]], [np.array(0.0), np.ones((5, 1))]) + np.testing.assert_allclose(out_vals[0].flat[0], 1.0, rtol=1e-12) + + def test_average_cell_centered_grid_length(self): + """``avg``'s length computation also has the cell-centered + (``len(coord) == values.shape[axis]``) extra-spacing branch.""" + x_cc = np.linspace(0.1, 0.9, 5) # total length 1.0 + values = 2.0 * np.ones((5, 1)) + _, out_vals = ev_ops.average([None, [x_cc]], [np.array(0.0), values]) + np.testing.assert_allclose(out_vals[0].flat[0], 2.0, rtol=1e-10) + + +class TestDivergence: + + def test_uniform_field_zero_divergence(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 1)) + _, out_vals = ev_ops.divergence([grid], [values]) + np.testing.assert_allclose(out_vals[0], 0.0, atol=1e-8) + + def test_linear_field_matches_analytic_divergence(self): + grid = [np.linspace(0.0, 1.0, 21)] + zc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = (2.0 * zc)[:, np.newaxis] # d/dx (2x) = 2 + _, out_vals = ev_ops.divergence([grid], [values]) + np.testing.assert_allclose(out_vals[0][:, 0], 2.0, rtol=1e-8) + + def test_too_many_components_raises(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 3)) # 3 comps, 1 dim + with pytest.raises(ValueError, match="longer than number of dimensions"): + ev_ops.divergence([grid], [values]) + + +class TestCurl: + + def test_1d_requires_3_components(self): + grid = [np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 2)) + with pytest.raises(ValueError, match="requires 3-component"): + ev_ops.curl([grid], [values]) + + def test_1d_curl_matches_analytic(self): + grid = [np.linspace(0.0, 1.0, 21)] + zc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = np.zeros((20, 3)) + values[:, 1] = zc # f_y = x -> curl_z = d(f_y)/dx = 1 + values[:, 2] = 2.0 * zc # f_z = 2x -> curl_y = -d(f_z)/dx = -2 + _, out_vals = ev_ops.curl([grid], [values]) + np.testing.assert_allclose(out_vals[0][:, 1], -2.0, rtol=1e-8) + np.testing.assert_allclose(out_vals[0][:, 2], 1.0, rtol=1e-8) + + def test_2d_too_few_components_raises(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 5, 1)) + with pytest.raises(ValueError, match="smaller than number of dimensions"): + ev_ops.curl([grid], [values]) + + def test_2d_exactly_two_components_computes_scalar_curl(self): + """The legacy code printed a misleading 'too long' WARNING for this + exact-match (num_comps == num_dims == 2) case and then computed the + standard 2D (in-plane) curl anyway -- a message bug, not a real + anomaly, fixed here by dropping the false-positive message. This is + the normal way to take the curl of a 2D vector field.""" + grid = [np.linspace(0.0, 1.0, 21), np.linspace(0.0, 1.0, 21)] + zc = 0.5 * (grid[0][:-1] + grid[0][1:]) + X, Y = np.meshgrid(zc, zc, indexing="ij") + values = np.zeros((20, 20, 2)) + values[..., 0] = -Y # f_x = -y + values[..., 1] = X # f_y = x -> curl_z = df_y/dx - df_x/dy = 1 - (-1) = 2 + _, out_vals = ev_ops.curl([grid], [values]) + assert out_vals[0].shape[-1] == 1 + np.testing.assert_allclose(out_vals[0][..., 0], 2.0, rtol=1e-6) + + def test_2d_three_components_computes_full_curl(self): + grid = [np.linspace(0.0, 1.0, 21), np.linspace(0.0, 1.0, 21)] + values = np.ones((20, 20, 3)) + _, out_vals = ev_ops.curl([grid], [values]) + assert out_vals[0].shape[-1] == 3 + + def test_2d_too_many_components_raises(self): + grid = [np.linspace(0.0, 1.0, 6), np.linspace(0.0, 1.0, 6)] + values = np.ones((5, 5, 4)) + with pytest.raises(ValueError, match="longer than number of dimensions"): + ev_ops.curl([grid], [values]) + + def test_3d_too_few_components_raises(self): + grid = [np.linspace(0.0, 1.0, 6)] * 3 + values = np.ones((5, 5, 5, 2)) + with pytest.raises(ValueError, match="smaller than number of dimensions"): + ev_ops.curl([grid], [values]) + + def test_3d_too_many_components_raises(self): + grid = [np.linspace(0.0, 1.0, 6)] * 3 + values = np.ones((5, 5, 5, 4)) + with pytest.raises(ValueError, match="longer than number of dimensions"): + ev_ops.curl([grid], [values]) + + def test_3d_curl_of_uniform_field_is_zero(self): + grid = [np.linspace(0.0, 1.0, 6)] * 3 + values = np.ones((5, 5, 5, 3)) + _, out_vals = ev_ops.curl([grid], [values]) + np.testing.assert_allclose(out_vals[0], 0.0, atol=1e-10) + + +class TestScaleComp: + + def test_slice_spec_scales_selected_components(self): + grid = [np.linspace(0.0, 1.0, 2)] + data = np.array([[1.0, 2.0, 3.0, 4.0]]) + out_grid, out_vals = ev_ops.scale_comp([None, None, grid], + [np.array(10.0), "1:3", data]) + np.testing.assert_allclose(out_vals[0], [[1.0, 20.0, 30.0, 4.0]]) + assert out_grid[0] is grid + + def test_int_array_spec(self): + data = np.array([[1.0, 2.0, 3.0]]) + _, out_vals = ev_ops.scale_comp( + [None, None, None], [np.array(2.0), np.array(1.0), data]) + np.testing.assert_allclose(out_vals[0], [[1.0, 4.0, 3.0]]) + + def test_bare_int_spec(self): + data = np.array([[1.0, 2.0, 3.0]]) + _, out_vals = ev_ops.scale_comp([None, None, None], + [np.array(5.0), 2, data]) + np.testing.assert_allclose(out_vals[0], [[1.0, 2.0, 15.0]]) + + def test_comma_list_spec(self): + data = np.array([[1.0, 2.0, 3.0]]) + _, out_vals = ev_ops.scale_comp([None, None, None], + [np.array(2.0), "0,2", data]) + np.testing.assert_allclose(out_vals[0], [[2.0, 2.0, 6.0]]) + + def test_does_not_mutate_original(self): + data = np.array([[1.0, 2.0, 3.0]]) + ev_ops.scale_comp([None, None, None], [np.array(10.0), "0:1", data]) + np.testing.assert_allclose(data, [[1.0, 2.0, 3.0]]) + + +class TestScaleZiAxis: + + def test_scales_named_axis(self): + axis = np.array([0.0, 1.0, 2.0]) + grid = [axis] + data = np.array([[1.0], [2.0], [3.0]]) + out_grid, out_vals = ev_ops.scale_zi_axis( + [None, None, grid], + [np.array(10.0), np.array(0.0), data]) + np.testing.assert_allclose(out_grid[0][0], [0.0, 10.0, 20.0]) + np.testing.assert_allclose(out_vals[0], data) diff --git a/tests/test_numerics_fft.py b/tests/test_numerics_fft.py new file mode 100644 index 00000000..b2ee15b2 --- /dev/null +++ b/tests/test_numerics_fft.py @@ -0,0 +1,371 @@ +"""Tests for postgkyl.numerics.fft -- fft/psd/iso and the polar helpers.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.numerics.fft import fft, init_polar, polar_isotropic + + +class TestFft1D: + + def test_returns_freq_and_ft_values(self): + N = 32 + grid = [np.linspace(0.0, 1.0, N + 1)] + x_cc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = np.sin(2 * np.pi * x_cc)[:, np.newaxis] + freq, ft = fft(grid, values) + assert len(freq) == 1 + assert ft.shape[0] == N + + def test_analytic_fft_of_pure_sine(self): + """A pure sine of frequency f0, sampled on a grid whose length matches + the values (``fft`` reads ``N``/``dx`` straight off ``len(grid[0])``, + so the grid array must already be the N-length sample-location axis, + not an N+1 nodal/edge array), has FFT power at exactly bins +-f0 and + zero elsewhere: an analytic, hand-computable reference.""" + N = 64 + x = np.linspace(0.0, 1.0, N, endpoint=False) + grid = [x] + f0 = 4 + values = np.sin(2 * np.pi * f0 * x)[:, np.newaxis] + freq, ft = fft(grid, values) + power = np.abs(ft[:, 0]) + i_pos = np.argmin(np.abs(freq[0] - f0)) + i_neg = np.argmin(np.abs(freq[0] + f0)) + np.testing.assert_allclose(freq[0][i_pos], f0, atol=1e-9) + np.testing.assert_allclose(freq[0][i_neg], -f0, atol=1e-9) + total = np.sum(power**2) + peak = power[i_pos]**2 + power[i_neg]**2 + assert peak / total > 0.999 + + def test_dc_component_for_constant(self): + N = 16 + grid = [np.linspace(0.0, 1.0, N + 1)] + values = np.ones((N, 1)) + freq, ft = fft(grid, values) + np.testing.assert_allclose(np.abs(ft[0, 0]), float(N)) + + def test_psd_halves_spectrum(self): + N = 32 + grid = [np.linspace(0.0, 1.0, N + 1)] + x_cc = 0.5 * (grid[0][:-1] + grid[0][1:]) + values = np.sin(2 * np.pi * x_cc)[:, np.newaxis] + freq, ft = fft(grid, values, psd=True) + assert ft.shape[0] == N // 2 + assert ft.shape[0] == len(freq[0]) + + def test_multiple_components(self): + N = 16 + grid = [np.linspace(0.0, 1.0, N + 1)] + values = np.column_stack([np.ones(N), np.zeros(N)]) + freq, ft = fft(grid, values) + assert ft.shape[-1] == 2 + + def test_dummy_dimension_squeezed(self): + N = 16 + grid = [np.linspace(0.0, 1.0, N + 1), np.array([0.0, 1.0])] + values = np.ones((N, 1, 1)) + freq, ft = fft(grid, values) + assert len(freq) == 1 + + +class TestFft2D: + + def test_2d_fft_returns_correct_shape(self): + Nx, Ny = 16, 8 + grid = [np.linspace(0.0, 1.0, Nx + 1), np.linspace(0.0, 1.0, Ny + 1)] + values = np.ones((Nx, Ny, 1)) + freq, ft = fft(grid, values) + assert ft.shape == (Nx, Ny, 1) + assert len(freq) == 2 + + def test_2d_psd(self): + Nx, Ny = 16, 8 + grid = [np.linspace(0.0, 1.0, Nx + 1), np.linspace(0.0, 1.0, Ny + 1)] + values = np.ones((Nx, Ny, 1)) + freq, ft = fft(grid, values, psd=True) + assert ft.shape[0] == Nx // 2 + assert ft.shape[1] == Ny // 2 + + def test_2d_psd_shape(self): + Nx, Ny = 8, 8 + grid = [np.linspace(0.0, 1.0, Nx + 1), np.linspace(0.0, 1.0, Ny + 1)] + values = np.ones((Nx, Ny, 1)) + freq, ft = fft(grid, values, psd=True) + assert ft.shape == (Nx // 2, Ny // 2, 1) + + +class TestFft3D: + + def test_3d_fft_runs(self): + Nx, Ny, Nz = 8, 8, 8 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + values = np.ones((Nx, Ny, Nz, 1)) + freq, ft = fft(grid, values) + assert ft.shape == (Nx, Ny, Nz, 1) + + def test_3d_psd_halves_dims(self): + Nx, Ny, Nz = 8, 8, 8 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + values = np.ones((Nx, Ny, Nz, 1)) + freq, ft = fft(grid, values, psd=True) + assert ft.shape == (Nx // 2, Ny // 2, Nz // 2, 1) + + def test_3d_psd_no_iso(self): + Nx, Ny, Nz = 4, 4, 4 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + rng = np.random.default_rng(0) + values = rng.random((Nx, Ny, Nz, 1)) + freq, ft = fft(grid, values, psd=True, iso=False) + assert ft.shape == (Nx // 2, Ny // 2, Nz // 2, 1) + + @pytest.mark.filterwarnings( + "ignore:invalid value encountered in divide:RuntimeWarning") + def test_3d_multi_comp(self): + Nx, Ny, Nz = 4, 4, 4 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + rng = np.random.default_rng(1) + values = rng.random((Nx, Ny, Nz, 3)) + freq, ft = fft(grid, values, psd=True, iso=True) + assert ft.shape[-1] == 3 + + +@pytest.mark.filterwarnings( + "ignore:invalid value encountered in divide:RuntimeWarning") +class TestFftIsotropic: + + def test_fft_3d_psd_iso(self): + Nx, Ny, Nz = 4, 4, 4 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + rng = np.random.default_rng(2) + values = rng.random((Nx, Ny, Nz, 1)) + freq, ft = fft(grid, values, psd=True, iso=True) + assert isinstance(freq, list) + assert len(freq) == 1 + assert ft.ndim == 2 + + def test_fft_3d_psd_iso_positive(self): + Nx, Ny, Nz = 4, 4, 4 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + values = np.ones((Nx, Ny, Nz, 1)) + freq, ft = fft(grid, values, psd=True, iso=True) + finite_vals = ft[np.isfinite(ft)] + assert np.all(finite_vals >= 0) + + def test_iso_preserves_total_power_end_to_end(self): + """Physically meaningful invariant that line coverage alone cannot see: + shell-averaging redistributes power onto k-shells but must not lose or + gain any of it. Reconstruct the Cartesian PSD independently + (iso=False) and check that weighting each isotropic bin by its shell's + cell count (``nbin``) reconstructs the same total.""" + Nx, Ny, Nz = 8, 8, 8 + grid = [ + np.linspace(0.0, 1.0, Nx + 1), + np.linspace(0.0, 1.0, Ny + 1), + np.linspace(0.0, 1.0, Nz + 1) + ] + rng = np.random.default_rng(4) + values = rng.random((Nx, Ny, Nz, 1)) + + freq_cart, ft_cartesian = fft(grid, values, psd=True, iso=False) + kx, ky, kz = freq_cart[0], freq_cart[1], freq_cart[2] + nkx, nky, nkz = len(kx), len(ky), len(kz) + # fft() derives nkpolar from the *nodal* grid lengths (Nx+1 here), not + # from the cell counts -- match that exactly to reproduce its binning. + N = np.array([len(grid[0]), len(grid[1]), len(grid[2])]) + nkpolar = int(np.sqrt(np.sum(N**2))) + _, nbin, polar_index, _ = init_polar(nkx, nky, nkz, kx, ky, kz, nkpolar) + expected_iso = polar_isotropic(nkpolar, nkx, nky, nkz, polar_index, nbin, + ft_cartesian[..., 0], kx, ky, kz) + + _, ft_iso = fft(grid, values, psd=True, iso=True) + + # fft(iso=True) must agree with a direct call to the same binning helpers. + np.testing.assert_allclose(ft_iso[:, 0], expected_iso) + + mask = nbin > 0 + total_from_shells = np.sum(ft_iso[mask, 0] * nbin[mask]) + np.testing.assert_allclose(total_from_shells, + np.sum(ft_cartesian[..., 0]), + rtol=1e-10) + + def test_iso_on_2d_data_treats_z_as_degenerate(self): + """iso doesn't check num_dims itself -- for 2D data the (dummy, unset) + third wavenumber axis is left at ``nkz=0``, so ``init_polar`` takes + its 2D branch and produces a 1D isotropic spectrum, same as 3D.""" + Nx, Ny = 8, 8 + grid = [np.linspace(0.0, 1.0, Nx + 1), np.linspace(0.0, 1.0, Ny + 1)] + values = np.ones((Nx, Ny, 1)) + freq, ft = fft(grid, values, psd=True, iso=True) + assert isinstance(freq, list) and len(freq) == 1 + assert ft.ndim == 2 + + +class TestFftPsdOnlySupported1D2D3D: + + def test_4d_raises_a_clean_value_error(self): + """src_bak's ``Only 1D, 2D, and 3D`` guard lived deep inside the psd + branch, behind a fixed-size ``N = np.zeros(3)`` that always raised a + confusing IndexError first for num_dims > 3 (psd or not) -- an + unreachable check. Fixed to raise the clean ValueError up front; every + working (<=3D) input is unaffected.""" + grid = [np.linspace(0.0, 1.0, 3)] * 4 + values = np.ones((2, 2, 2, 2, 1)) + with pytest.raises(ValueError, match="1D, 2D, and 3D"): + fft(grid, values) + with pytest.raises(ValueError, match="1D, 2D, and 3D"): + fft(grid, values, psd=True) + + +# --------------------------------------------------------------------------- +# init_polar +# --------------------------------------------------------------------------- + + +class TestInitPolar: + + def test_nkpolar_zero_returns_empty(self): + akp, nbin, polar_index, akplim = init_polar(4, 4, 0, [], [], [], 0) + assert akp == [] + assert nbin == 0 + assert polar_index == [] + assert akplim == [] + + def test_2d_case_basic(self): + N = 8 + kx = np.fft.fftfreq(N, 1.0 / N)[:N // 2] + ky = np.fft.fftfreq(N, 1.0 / N)[:N // 2] + nkpolar = 5 + akp, nbin, polar_index, akplim = init_polar(len(kx), len(ky), 0, kx, ky, [], + nkpolar) + assert len(akp) == nkpolar + assert len(nbin) == nkpolar + assert polar_index.shape == (len(kx), len(ky)) + assert len(akplim) == nkpolar + 1 + assert np.sum(nbin) > 0 + + def test_2d_case_nkx1(self): + kx = np.array([0.0]) + ky = np.array([0.0, 1.0, 2.0]) + akp, nbin, polar_index, akplim = init_polar(1, 3, 0, kx, ky, [], 3) + assert len(akp) == 3 + + def test_2d_case_nky1(self): + kx = np.array([0.0, 1.0, 2.0]) + ky = np.array([0.0]) + akp, nbin, polar_index, akplim = init_polar(3, 1, 0, kx, ky, [], 3) + assert len(akp) == 3 + + def test_2d_case_nkx1_and_nky1_uses_zero_spacing(self): + """The ``nkx == 1 and nky == 1`` branch (dkp = 0): a single-cell grid + in both directions. Also proves the fixed ``and`` (src_bak's ``&`` + precedence bug would have made the parity of nky, not the actual + nkx/nky == 1 check, decide this branch).""" + kx = np.array([0.0]) + ky = np.array([0.0]) + akp, nbin, polar_index, akplim = init_polar(1, 1, 0, kx, ky, [], 2) + np.testing.assert_allclose(akp, [0.0, 0.0]) + + def test_3d_case_basic(self): + N = 4 + kx = np.fft.fftfreq(N)[:N // 2] + ky = np.fft.fftfreq(N)[:N // 2] + kz = np.fft.fftfreq(N)[:N // 2] + nkpolar = 4 + akp, nbin, polar_index, akplim = init_polar(len(kx), len(ky), len(kz), kx, + ky, kz, nkpolar) + assert len(akp) == nkpolar + assert polar_index.shape == (len(kx), len(ky), len(kz)) + assert np.sum(nbin) > 0 + + def test_3d_case_nkx1(self): + kx = np.array([0.0]) + ky = np.array([0.0, 1.0]) + kz = np.array([0.0, 1.0]) + akp, nbin, polar_index, akplim = init_polar(1, 2, 2, kx, ky, kz, 2) + assert len(akp) == 2 + + def test_3d_case_nky1(self): + kx = np.array([0.0, 1.0]) + ky = np.array([0.0]) + kz = np.array([0.0, 1.0]) + akp, nbin, polar_index, akplim = init_polar(2, 1, 2, kx, ky, kz, 2) + assert len(akp) == 2 + + def test_3d_case_nkz1(self): + kx = np.array([0.0, 1.0]) + ky = np.array([0.0, 1.0]) + kz = np.array([0.0]) + akp, nbin, polar_index, akplim = init_polar(2, 2, 1, kx, ky, kz, 2) + assert len(akp) == 2 + + def test_3d_case_all_singleton_uses_zero_spacing(self): + """The ``nkx == 1 and nky == 1 and nkz == 1`` branch (dkp = 0), and a + parity-sensitive proof of the fixed ``and`` (an even count anywhere + would have flipped src_bak's ``&``-precedence-bugged condition).""" + kx = ky = kz = np.array([0.0]) + akp, nbin, polar_index, akplim = init_polar(1, 1, 1, kx, ky, kz, 2) + np.testing.assert_allclose(akp, [0.0, 0.0]) + + +# --------------------------------------------------------------------------- +# polar_isotropic +# --------------------------------------------------------------------------- + + +class TestPolarIsotropic: + + def test_2d_case(self): + N = 8 + kx = np.fft.fftfreq(N)[:N // 2] + ky = np.fft.fftfreq(N)[:N // 2] + nkpolar = 3 + akp, nbin, polar_index, _ = init_polar(len(kx), len(ky), 0, kx, ky, [], + nkpolar) + fft_matrix = np.ones((len(kx), len(ky))) + result = polar_isotropic(nkpolar, len(kx), len(ky), 0, polar_index, nbin, + fft_matrix, kx, ky, []) + assert result.shape == (nkpolar, ) + assert np.any(nbin > 0) + + @pytest.mark.filterwarnings( + "ignore:invalid value encountered in divide:RuntimeWarning") + def test_3d_case(self): + N = 4 + kx = np.fft.fftfreq(N)[:N // 2] + ky = np.fft.fftfreq(N)[:N // 2] + kz = np.fft.fftfreq(N)[:N // 2] + nkpolar = 3 + akp, nbin, polar_index, _ = init_polar(len(kx), len(ky), len(kz), kx, ky, + kz, nkpolar) + fft_matrix = np.ones((len(kx), len(ky), len(kz))) + result = polar_isotropic(nkpolar, len(kx), len(ky), len(kz), polar_index, + nbin, fft_matrix, kx, ky, kz) + assert result.shape == (nkpolar, ) + assert np.any(nbin > 0) diff --git a/tests/test_numerics_filters.py b/tests/test_numerics_filters.py new file mode 100644 index 00000000..87d294ef --- /dev/null +++ b/tests/test_numerics_filters.py @@ -0,0 +1,86 @@ +"""Tests for postgkyl.numerics.filters -- fft_filtering and butter_filtering.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.numerics.filters import fft_filtering, butter_filtering + + +class TestFftFiltering: + + def test_removes_high_frequency_component(self): + N = 256 + dt = 1.0 / N + t = np.linspace(0.0, 1.0 - dt, N) + signal = np.sin(2 * 2 * np.pi * t) + 0.5 * np.sin(50 * 2 * np.pi * t) + filtered = fft_filtering(signal, dt=dt, cutoff=10.0) + high_freq_power_before = 0.5 + high_freq_power_after = np.std( + np.real(filtered) - np.sin(2 * 2 * np.pi * t)) + assert high_freq_power_after < 0.1 * high_freq_power_before + + def test_preserves_dc_component(self): + N = 128 + dt = 1.0 / N + signal = np.ones(N) * 3.0 + filtered = fft_filtering(signal, dt=dt, cutoff=1.0) + np.testing.assert_allclose(np.real(filtered), 3.0, atol=1e-10) + + def test_output_same_length(self): + N = 64 + rng = np.random.default_rng(0) + signal = rng.standard_normal(N) + filtered = fft_filtering(signal, dt=0.01, cutoff=10.0) + assert len(filtered) == N + + def test_cutoff_zero_removes_all(self): + N = 64 + signal = np.sin(2 * np.pi * np.linspace(0, 1, N)) + filtered = fft_filtering(signal, dt=1.0 / N, cutoff=0.0) + np.testing.assert_allclose(np.abs(filtered).max(), 0.0, atol=1e-10) + + def test_cutoff_is_keyword_only(self): + with pytest.raises(TypeError): + fft_filtering(np.ones(8), 1.0, 5.0) # type: ignore[misc] + + def test_cutoff_required(self): + with pytest.raises(TypeError): + fft_filtering(np.ones(8)) # type: ignore[call-arg] + + +class TestButterFiltering: + + def test_removes_high_frequency(self): + N = 512 + dt = 1.0 / N + t = np.linspace(0.0, 1.0, N) + low = np.sin(2 * 2 * np.pi * t) + high = 0.5 * np.sin(100 * 2 * np.pi * t) + filtered = butter_filtering(low + high, dt=dt, cutoff=10.0) + skip = N // 5 + std_filtered = np.std(filtered[skip:]) + std_original = np.std((low + high)[skip:]) + assert std_filtered < std_original + + def test_output_same_length(self): + N = 64 + rng = np.random.default_rng(1) + signal = rng.standard_normal(N) + filtered = butter_filtering(signal, dt=0.01, cutoff=5.0) + assert len(filtered) == N + + def test_preserves_low_frequency(self): + N = 512 + dt = 1.0 / N + t = np.linspace(0.0, 1.0, N) + freq = 1.0 + signal = np.sin(2 * np.pi * freq * t) + filtered = butter_filtering(signal, dt=dt, cutoff=100.0) + skip = N // 5 + np.testing.assert_allclose(np.max(np.abs(filtered[skip:])), 1.0, atol=0.05) + + def test_cutoff_required(self): + with pytest.raises(TypeError): + butter_filtering(np.ones(8)) # type: ignore[call-arg] diff --git a/tests/test_numerics_fit.py b/tests/test_numerics_fit.py new file mode 100644 index 00000000..69e2aba9 --- /dev/null +++ b/tests/test_numerics_fit.py @@ -0,0 +1,566 @@ +"""Tests for postgkyl.numerics.fit -- model functions, RPN parser, fit/auto_guess. + +Ports the array-only subset of ``tests_bak/test_fit.py``: the model +functions, ``fit``/``fit_evaluate``, and the RPN expression machinery. +``FitTypeParam`` and the ``fit`` CLI command belong to the CLI/operations layers +and are not part of this leaf module -- they are not ported here. +""" + +from __future__ import annotations + +import importlib + +import numpy as np +import pytest + +# `postgkyl.numerics.fit` (the submodule) is shadowed by the `fit` FUNCTION +# that numerics/__init__.py re-exports under the same attribute name -- see +# the note in tests/test_coverage_leaf.py. importlib sidesteps the +# package's __init__ entirely and returns the actual submodule object. +fitmod = importlib.import_module("postgkyl.numerics.fit") + +# ── model functions ────────────────────────────────────────────────────────── + + +class TestFitFunctions: + + def test_linear_evaluation(self): + x = np.array([0.0, 1.0, 2.0]) + np.testing.assert_allclose(fitmod.linear(x, 3.0, -1.0), [-1.0, 2.0, 5.0]) + + def test_quadratic_evaluation(self): + x = np.array([0.0, 1.0, 2.0, 3.0]) + np.testing.assert_allclose(fitmod.quadratic(x, 1.0, -2.0, 1.0), + [1.0, 0.0, 1.0, 4.0]) + + def test_plane_evaluation(self): + XY = np.array([[0.0, 1.0], [0.0, 1.0]]) + np.testing.assert_allclose(fitmod.plane(XY, 2.0, -1.0, 0.5), [0.5, 1.5]) + + def test_quadratic2d_evaluation(self): + XY = np.array([[1.0], [2.0]]) + result = fitmod.quadratic2d(XY, 1.0, 0.0, 0.0, 0.0, 0.0, 3.0) + np.testing.assert_allclose(result, [4.0]) + + def test_exp_plateau_evaluation(self): + x = np.array([0.0, 1.0]) + np.testing.assert_allclose(fitmod.exp_plateau(x, 2.0, 0.0, 1.0), [3.0, 3.0]) + + def test_gaussian_evaluation(self): + x = np.array([0.0]) + np.testing.assert_allclose(fitmod.gaussian(x, 3.0, 0.0, 1.0), [3.0]) + + def test_power_evaluation(self): + x = np.array([1.0, 2.0, 4.0]) + np.testing.assert_allclose(fitmod.power(x, 2.0, 3.0, 1.0), + [3.0, 17.0, 129.0]) + + def test_sinusoid_evaluation(self): + x = np.array([0.0, np.pi / 2]) + np.testing.assert_allclose(fitmod.sinusoid(x, 1.0, 1.0, 0.0, 0.5), + [0.5, 1.5], + atol=1e-14) + + def test_tanh_transition_evaluation(self): + x = np.array([0.0]) + np.testing.assert_allclose(fitmod.tanh_transition(x, 2.0, 0.0, 1.0, -1.0), + [-1.0]) + + def test_exp2_evaluation(self): + np.testing.assert_allclose(fitmod.exp2(0.0, a=2.0, b=1.0), 2.0) + x = np.array([0.0, 1.0, 2.0]) + np.testing.assert_allclose(fitmod.exp2(x, a=1.0, b=1.0), np.exp(2 * x)) + + def test_fit_functions_and_ndim_consistent(self): + assert set(fitmod.FIT_FUNCTIONS) == set(fitmod.FIT_NDIM) + + def test_fit_ndim_values(self): + assert fitmod.FIT_NDIM["linear"] == 1 + assert fitmod.FIT_NDIM["quadratic"] == 1 + assert fitmod.FIT_NDIM["plane"] == 2 + assert fitmod.FIT_NDIM["quadratic2d"] == 2 + assert fitmod.FIT_NDIM["exp_plateau"] == 1 + assert fitmod.FIT_NDIM["gaussian"] == 1 + assert fitmod.FIT_NDIM["power"] == 1 + assert fitmod.FIT_NDIM["sinusoid"] == 1 + assert fitmod.FIT_NDIM["tanh_transition"] == 1 + assert fitmod.FIT_NDIM["exp2"] == 1 + + def test_fit_evaluate_builtin(self): + x = np.array([0.0, 1.0, 2.0]) + out = fitmod.fit_evaluate(x, "linear", [3.0, -1.0]) + np.testing.assert_allclose(out, [-1.0, 2.0, 5.0]) + + def test_fit_evaluate_rpn(self): + x = np.array([0.0, 1.0, 2.0]) + out = fitmod.fit_evaluate(x, "a x * b +", [3.0, -1.0]) + np.testing.assert_allclose(out, [-1.0, 2.0, 5.0]) + + +# ── fit() -- 1-D models ────────────────────────────────────────────────────── + + +class TestFit1D: + + def test_linear_exact_data_recovers_params(self): + x = np.linspace(0, 10, 50) + y = 3.0 * x - 1.5 + params, _, R2 = fitmod.fit(x, y, "linear") + np.testing.assert_allclose(params, [3.0, -1.5], rtol=1e-10) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_quadratic_exact_data_recovers_params(self): + x = np.linspace(-2, 2, 60) + y = 0.5 * x**2 - 1.0 * x + 2.0 + params, _, R2 = fitmod.fit(x, y, "quadratic") + np.testing.assert_allclose(params, [0.5, -1.0, 2.0], rtol=1e-10) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_linear_noisy_data_high_R2_and_close_params(self): + rng = np.random.default_rng(0) + x = np.linspace(0, 10, 200) + y = 2.0 * x + 1.0 + rng.normal(0, 0.1, 200) + params, _, R2 = fitmod.fit(x, y, "linear") + assert R2 > 0.999 + np.testing.assert_allclose(params[0], 2.0, atol=0.05) + np.testing.assert_allclose(params[1], 1.0, atol=0.1) + + def test_returns_covariance_with_correct_shape(self): + x = np.linspace(0, 5, 30) + y = x + 1.0 + _, cov, _ = fitmod.fit(x, y, "linear") + assert cov.shape == (2, 2) + + def test_initial_guess_does_not_change_result_on_exact_data(self): + x = np.linspace(0, 10, 50) + y = 5.0 * x + 3.0 + params_default, _, _ = fitmod.fit(x, y, "linear") + params_guess, _, _ = fitmod.fit(x, y, "linear", p0=[10.0, 10.0]) + np.testing.assert_allclose(params_default, params_guess, rtol=1e-8) + + def test_exp_plateau_exact_data_recovers_params(self): + x = np.linspace(0, 5, 80) + true_params = [3.0, -1.5, 1.0] + y = fitmod.exp_plateau(x, *true_params) + params, _, R2 = fitmod.fit(x, y, "exp_plateau", p0=[1.0, -1.0, 0.0]) + np.testing.assert_allclose(params, true_params, rtol=1e-6) + assert R2 == pytest.approx(1.0, abs=1e-8) + + def test_exp_plateau_noisy_data_high_R2(self): + rng = np.random.default_rng(7) + x = np.linspace(0, 5, 100) + y = fitmod.exp_plateau(x, 3.0, -1.5, 1.0) + rng.normal(0, 0.05, 100) + _, _, R2 = fitmod.fit(x, y, "exp_plateau", p0=[1.0, -1.0, 0.0]) + assert R2 > 0.99 + + def test_invalid_fit_type_raises_value_error(self): + x = np.linspace(0, 1, 10) + y = x + with pytest.raises(ValueError, match="not recognized"): + fitmod.fit(x, y, "cubic") + + def test_gaussian_exact_data_recovers_params(self): + x = np.linspace(-3, 3, 100) + true_params = [2.0, 0.5, 0.8] + y = fitmod.gaussian(x, *true_params) + params, _, R2 = fitmod.fit(x, y, "gaussian", p0=[1.0, 0.0, 1.0]) + np.testing.assert_allclose(params, true_params, rtol=1e-6) + assert R2 == pytest.approx(1.0, abs=1e-8) + + def test_power_exact_data_recovers_params(self): + x = np.linspace(1, 5, 60) + true_params = [3.0, 2.0, -1.0] + y = fitmod.power(x, *true_params) + params, _, R2 = fitmod.fit(x, y, "power", p0=[1.0, 1.5, 0.0]) + np.testing.assert_allclose(params, true_params, rtol=1e-6) + assert R2 == pytest.approx(1.0, abs=1e-8) + + def test_sinusoid_exact_data_recovers_params(self): + x = np.linspace(0, 4 * np.pi, 200) + true_params = [2.0, 1.0, 0.3, 0.5] + y = fitmod.sinusoid(x, *true_params) + params, _, R2 = fitmod.fit(x, y, "sinusoid", p0=[1.5, 1.0, 0.0, 0.0]) + np.testing.assert_allclose(params, true_params, rtol=1e-5) + assert R2 == pytest.approx(1.0, abs=1e-8) + + def test_tanh_transition_exact_data_recovers_params(self): + x = np.linspace(-5, 5, 100) + true_params = [3.0, 1.0, 0.5, 2.0] + y = fitmod.tanh_transition(x, *true_params) + params, _, R2 = fitmod.fit(x, y, "tanh_transition", p0=[1.0, 0.0, 1.0, 0.0]) + np.testing.assert_allclose(params, true_params, rtol=1e-6) + assert R2 == pytest.approx(1.0, abs=1e-8) + + +# ── RPN expression support ─────────────────────────────────────────────────── + + +class TestRPN: + + def test_param_names_basic(self): + assert fitmod.rpn_param_names("a x * b +") == ["a", "b"] + + def test_param_names_excludes_spatial_vars(self): + assert "x" not in fitmod.rpn_param_names("a x * b +") + assert "y" not in fitmod.rpn_param_names("a x * b y * + c +") + + def test_param_names_excludes_operators(self): + assert "+" not in fitmod.rpn_param_names("a x * b +") + assert "*" not in fitmod.rpn_param_names("a x * b +") + + def test_param_names_excludes_functions(self): + assert "exp" not in fitmod.rpn_param_names("A b x * exp *") + + def test_param_names_excludes_numeric_literals(self): + assert fitmod.rpn_param_names("2 x * 1 +") == [] + + def test_param_names_preserves_order(self): + assert fitmod.rpn_param_names("A b x * exp * C +") == ["A", "b", "C"] + + def test_param_names_empty_expression(self): + assert fitmod.rpn_param_names("") == [] + + def test_ndim_1d(self): + assert fitmod.rpn_ndim("a x * b +") == 1 + + def test_ndim_2d(self): + assert fitmod.rpn_ndim("a x * b y * + c +") == 2 + + def test_rpn_linear_recovers_params(self): + x = np.linspace(0, 10, 50) + y = 3.0 * x - 1.5 + params, _, R2 = fitmod.fit(x, y, "a x * b +", p0=[1.0, 0.0]) + np.testing.assert_allclose(params, [3.0, -1.5], rtol=1e-8) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_rpn_exp_recovers_params(self): + x = np.linspace(0, 3, 80) + true_A, true_b = 2.0, -0.5 + y = true_A * np.exp(true_b * x) + params, _, R2 = fitmod.fit(x, y, "A b x * exp *", p0=[1.0, -1.0]) + np.testing.assert_allclose(params, [true_A, true_b], rtol=1e-6) + assert R2 == pytest.approx(1.0, abs=1e-8) + + def test_rpn_plane_2d_recovers_params(self): + X, Y = np.meshgrid(np.linspace(0, 5, 15), + np.linspace(0, 3, 10), + indexing="ij") + xdata = np.array([X.flatten(), Y.flatten()]) + y = 2.0 * X.flatten() - 1.5 * Y.flatten() + 0.5 + params, _, R2 = fitmod.fit(xdata, + y, + "a x * b y * + c +", + p0=[1.0, 1.0, 0.0]) + np.testing.assert_allclose(params, [2.0, -1.5, 0.5], rtol=1e-8) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_rpn_literal_coefficients(self): + x = np.linspace(1, 5, 40) + y = 2.0 * x**2 + params, _, R2 = fitmod.fit(x, y, "a x 2 ** *", p0=[1.0]) + np.testing.assert_allclose(params, [2.0], rtol=1e-8) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_rpn_caret_power_operator(self): + x = np.linspace(1, 5, 40) + y = 2.0 * x**2 + params, _, R2 = fitmod.fit(x, y, "a x 2 ^ *", p0=[1.0]) + np.testing.assert_allclose(params, [2.0], rtol=1e-8) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_rpn_subtract_operator(self): + func = fitmod._rpn_make_func("x a -") + x = np.array([5.0, 10.0]) + np.testing.assert_allclose(func(x, 2.0), [3.0, 8.0]) + + def test_rpn_divide_operator(self): + func = fitmod._rpn_make_func("x a /") + x = np.array([10.0, 20.0]) + np.testing.assert_allclose(func(x, 2.0), [5.0, 10.0]) + + def test_rpn_pure_constant_expression_broadcasts_to_x_shape(self): + """A scalar-only RPN expression (no free params, no spatial var used) + still broadcasts its result to xdata's shape, since ``x`` is always + bound in the evaluation namespace regardless of whether the + expression actually references it.""" + func = fitmod._rpn_make_func("2 3 +") + x = np.array([0.0, 1.0, 2.0]) + out = func(x) + np.testing.assert_allclose(out, [5.0, 5.0, 5.0]) + + def test_rpn_malformed_stack_raises(self): + # Leading operator with empty stack causes IndexError inside curve_fit. + x = np.linspace(0, 1, 10) + with pytest.raises((IndexError, Exception)): + fitmod.fit(x, x, "* x a +", p0=[1.0]) + + def test_rpn_bad_token_raises_value_error(self): + """A token that is neither an operator, function, known parameter, nor + a valid float literal (bad token edge case).""" + x = np.linspace(0, 1, 10) + func = fitmod._rpn_make_func("a x * not_a_number +") + with pytest.raises(ValueError): + func(x, 1.0) + + def test_rpn_arity_mismatch_raises(self): + """Requesting fewer parameter values than the expression's free + parameters is an arity mismatch: the dict(zip(...)) call silently + drops the excess names, so the *unbound* stray name looks up as + missing from ``ns`` and falls through to ``float(tok)``, which raises + ValueError on a non-numeric token.""" + func = fitmod._rpn_make_func("a b + x *") + with pytest.raises(ValueError): + func(np.array([1.0, 2.0]), 1.0) # only 'a' bound, 'b' unresolved + + def test_fittype_param_accepts_rpn_via_fit(self): + x = np.linspace(0, 10, 50) + y = 3.0 * x - 1.5 + params, _, _ = fitmod.fit(x, y, "a x * b +", p0=[1.0, 0.0]) + assert len(params) == 2 + + +# ── fit() -- 2-D models ────────────────────────────────────────────────────── + + +class TestFit2D: + + @staticmethod + def _xdata(x, y): + X, Y = np.meshgrid(x, y, indexing="ij") + return np.array([X.flatten(), Y.flatten()]) + + def test_plane_exact_data_recovers_params(self): + xdata = self._xdata(np.linspace(0, 5, 20), np.linspace(0, 3, 15)) + zdata = fitmod.plane(xdata, 2.0, -1.5, 0.5) + params, _, R2 = fitmod.fit(xdata, zdata, "plane") + np.testing.assert_allclose(params, [2.0, -1.5, 0.5], rtol=1e-10) + assert R2 == pytest.approx(1.0, abs=1e-10) + + def test_quadratic2d_exact_data_recovers_params(self): + xdata = self._xdata(np.linspace(0, 4, 15), np.linspace(0, 3, 12)) + true_params = [0.3, 0.2, -0.1, 1.0, -0.5, 2.0] + zdata = fitmod.quadratic2d(xdata, *true_params) + params, _, R2 = fitmod.fit(xdata, zdata, "quadratic2d") + np.testing.assert_allclose(params, true_params, rtol=1e-8) + assert R2 == pytest.approx(1.0, abs=1e-8) + + def test_plane_noisy_data_high_R2(self): + rng = np.random.default_rng(42) + xdata = self._xdata(np.linspace(0, 5, 30), np.linspace(0, 3, 25)) + zdata = fitmod.plane(xdata, 2.0, -1.5, 0.5) + rng.normal( + 0, 0.05, xdata.shape[1]) + _, _, R2 = fitmod.fit(xdata, zdata, "plane") + assert R2 > 0.999 + + def test_plane_returns_correct_covariance_shape(self): + xdata = self._xdata(np.linspace(0, 5, 10), np.linspace(0, 3, 8)) + zdata = fitmod.plane(xdata, 1.0, 2.0, 0.0) + _, cov, _ = fitmod.fit(xdata, zdata, "plane") + assert cov.shape == (3, 3) + + +# ── auto_guess ──────────────────────────────────────────────────────────────── + + +class TestAutoGuess: + + def test_returns_none_for_all_nan(self): + x = np.linspace(0, 1, 10) + y = np.full(10, np.nan) + assert fitmod.auto_guess("linear", x, y) is None + + def test_returns_none_for_rpn_expression(self): + x = np.linspace(0, 1, 10) + y = x + assert fitmod.auto_guess("a x * b +", x, y) is None + + def test_linear_guess_is_reasonable(self): + x = np.linspace(0, 10, 50) + y = 2.0 * x + 1.0 + a, b = fitmod.auto_guess("linear", x, y) + np.testing.assert_allclose([a, b], [2.0, 1.0], rtol=1e-6) + + def test_quadratic_guess_is_reasonable(self): + x = np.linspace(-2, 2, 60) + y = 0.5 * x**2 - 1.0 * x + 2.0 + guess = fitmod.auto_guess("quadratic", x, y) + np.testing.assert_allclose(guess, [0.5, -1.0, 2.0], rtol=1e-6) + + def test_quadratic_guess_falls_back_when_polyfit_raises(self): + """np.polyfit raises on an empty vector; auto_guess catches it and + falls back to a fixed placeholder guess rather than propagating.""" + x = np.array([]) + y = np.array([1.0, 2.0]) # not all-NaN, so the finite-check passes + guess = fitmod.auto_guess("quadratic", x, y) + assert guess == [0.0, 1.0, pytest.approx(1.5)] + + def test_plane_guess_is_reasonable(self): + X, Y = np.meshgrid(np.linspace(0, 5, 20), + np.linspace(0, 3, 15), + indexing="ij") + xdata = np.array([X.flatten(), Y.flatten()]) + y = 2.0 * X.flatten() - 1.5 * Y.flatten() + 0.5 + guess = fitmod.auto_guess("plane", xdata, y) + np.testing.assert_allclose(guess, [2.0, -1.5, 0.5], rtol=1e-6) + + def test_quadratic2d_guess_is_reasonable(self): + X, Y = np.meshgrid(np.linspace(0, 4, 15), + np.linspace(0, 3, 12), + indexing="ij") + xdata = np.array([X.flatten(), Y.flatten()]) + true_params = [0.3, 0.2, -0.1, 1.0, -0.5, 2.0] + y = fitmod.quadratic2d(xdata, *true_params) + guess = fitmod.auto_guess("quadratic2d", xdata, y) + np.testing.assert_allclose(guess, true_params, rtol=1e-6) + + def test_exp_plateau_guess_seeds_a_working_fit(self): + x = np.linspace(0, 5, 80) + true_params = [3.0, -1.5, 1.0] + y = fitmod.exp_plateau(x, *true_params) + guess = fitmod.auto_guess("exp_plateau", x, y) + params, _, R2 = fitmod.fit(x, y, "exp_plateau", p0=guess) + np.testing.assert_allclose(params, true_params, rtol=1e-4) + + def test_gaussian_guess_seeds_a_working_fit(self): + x = np.linspace(-3, 3, 100) + true_params = [2.0, 0.5, 0.8] + y = fitmod.gaussian(x, *true_params) + guess = fitmod.auto_guess("gaussian", x, y) + params, _, R2 = fitmod.fit(x, y, "gaussian", p0=guess) + np.testing.assert_allclose(params, true_params, rtol=1e-4) + + def test_gaussian_guess_narrow_peak_uses_fallback_sigma(self): + """When fewer than two points reach half-max, the FWHM estimate falls + back to a quarter of the domain width.""" + x = np.linspace(-3, 3, 7) + y = np.zeros_like(x) + y[3] = 5.0 # single spike -> only one point at/above half-max + guess = fitmod.auto_guess("gaussian", x, y) + assert guess[2] == pytest.approx((x.max() - x.min()) / 4) + + def test_power_guess_seeds_a_working_fit(self): + x = np.linspace(1, 5, 60) + true_params = [3.0, 2.0, -1.0] + y = fitmod.power(x, *true_params) + guess = fitmod.auto_guess("power", x, y) + params, _, R2 = fitmod.fit(x, y, "power", p0=guess) + np.testing.assert_allclose(params, true_params, rtol=1e-4) + + def test_sinusoid_guess_seeds_a_working_fit(self): + x = np.linspace(0, 4 * np.pi, 200) + true_params = [2.0, 1.0, 0.3, 0.5] + y = fitmod.sinusoid(x, *true_params) + guess = fitmod.auto_guess("sinusoid", x, y) + params, _, R2 = fitmod.fit(x, y, "sinusoid", p0=guess) + assert R2 > 0.99 + + def test_sinusoid_guess_single_point_omega_fallback(self): + x = np.array([0.0]) + y = np.array([1.0]) + guess = fitmod.auto_guess("sinusoid", x, y) + assert guess[1] == 1.0 # omega fallback for len(x) <= 1 + + def test_tanh_transition_guess_seeds_a_working_fit(self): + x = np.linspace(-5, 5, 100) + true_params = [3.0, 1.0, 0.5, 2.0] + y = fitmod.tanh_transition(x, *true_params) + guess = fitmod.auto_guess("tanh_transition", x, y) + params, _, R2 = fitmod.fit(x, y, "tanh_transition", p0=guess) + np.testing.assert_allclose(params, true_params, rtol=1e-4) + + def test_unknown_fit_type_returns_none(self): + x = np.linspace(0, 1, 10) + y = x + assert fitmod.auto_guess("not_a_real_model", x, y) is None + + def test_exp2_guess_seeds_a_working_fit(self): + x = np.linspace(0, 5, 80) + true_params = [1.0, 0.8] + y = fitmod.exp2(x, *true_params) + guess = fitmod.auto_guess("exp2", x, y) + params, _, R2 = fitmod.fit(x, y, "exp2", p0=guess) + np.testing.assert_allclose(params, true_params, rtol=1e-4) + + def test_exp2_guess_is_scale_invariant(self): + """The log-linear guess should converge without needing x normalized + to O(1) -- unlike a blind (1, 1) seed, it stays accurate as the time + axis grows.""" + x = np.linspace(0, 500, 200) + true_params = [2.0, 0.01] + y = fitmod.exp2(x, *true_params) + guess = fitmod.auto_guess("exp2", x, y) + params, _, R2 = fitmod.fit(x, y, "exp2", p0=guess) + np.testing.assert_allclose(params, true_params, rtol=1e-4) + + +# ── fit_best_window ─────────────────────────────────────────────────────────── + + +class TestFitBestWindow: + + def test_recovers_known_growth_rate(self): + x = np.linspace(0, 5, 60) + true_a, true_b = 1.0, 0.8 + y = fitmod.exp2(x, true_a, true_b) + params, cov, R2, n = fitmod.fit_best_window(x, y, "exp2") + assert R2 > 0.99 + np.testing.assert_allclose(params[1], true_b, rtol=0.05) + + def test_returns_four_elements(self): + x = np.linspace(0, 3, 30) + y = fitmod.exp2(x, 1.0, 0.5) + result = fitmod.fit_best_window(x, y, "exp2") + assert len(result) == 4 + + def test_best_n_is_within_bounds(self): + x = np.linspace(0, 4, 40) + y = fitmod.exp2(x, 1.0, 0.5) + _, _, _, n = fitmod.fit_best_window(x, y, "exp2", min_n=5) + assert 5 <= n <= len(x) + + def test_custom_min_n(self): + x = np.linspace(0, 3, 30) + y = fitmod.exp2(x, 1.0, 0.5) + _, _, _, n = fitmod.fit_best_window(x, y, "exp2", min_n=10) + assert n >= 10 + + def test_curve_fit_failure_for_some_windows_is_skipped(self, monkeypatch): + """A RuntimeError from curve_fit (non-convergence) for one fitting + window is caught, not fatal -- the scan continues and still returns + the best window that did converge.""" + x = np.linspace(0, 5, 30) + y = fitmod.exp2(x, 1.0, 0.8) + real_curve_fit = fitmod.opt.curve_fit + calls = {"n": 0} + + def flaky_curve_fit(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("simulated non-convergence") + return real_curve_fit(*args, **kwargs) + + monkeypatch.setattr(fitmod.opt, "curve_fit", flaky_curve_fit) + _, _, R2, _ = fitmod.fit_best_window(x, y, "exp2", min_n=5) + assert R2 > 0.9 + + def test_all_windows_failing_to_converge_raises(self, monkeypatch): + """If curve_fit never converges for any window in the scan range, + fit_best_window must raise a clear domain error rather than crash.""" + x = np.linspace(0, 5, 30) + y = fitmod.exp2(x, 1.0, 0.8) + + def always_fails(*args, **kwargs): + raise RuntimeError("simulated non-convergence") + + monkeypatch.setattr(fitmod.opt, "curve_fit", always_fails) + with pytest.raises(RuntimeError, match="failed to converge"): + fitmod.fit_best_window(x, y, "exp2", min_n=5) + + def test_generic_over_fit_type(self): + """fit_best_window is generic over any registered fit_type, not + hard-wired to exp2 -- generalizing the old growth-specific scan.""" + x = np.linspace(0.1, 5, 40) + y = 2.0 * x + 1.0 + params, _, R2, _ = fitmod.fit_best_window(x, y, "linear", p0=[1.0, 1.0]) + assert R2 > 0.99 + np.testing.assert_allclose(params, [2.0, 1.0], rtol=1e-6) diff --git a/tests/test_numerics_grid_centering.py b/tests/test_numerics_grid_centering.py new file mode 100644 index 00000000..b350ff1d --- /dev/null +++ b/tests/test_numerics_grid_centering.py @@ -0,0 +1,104 @@ +"""Tests for postgkyl.numerics.grid_centering -- nodal_to_cell_centered_grid.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl.numerics.grid_centering import nodal_to_cell_centered_grid + + +class TestNodalToCellCenteredGrid: + + def test_1d_nodal_grid_is_centered(self): + grid = [np.linspace(0.0, 1.0, 5)] # 4 cells, nodal (5 points) + out = nodal_to_cell_centered_grid(grid, cells=np.array([4])) + assert len(out) == 1 + assert out[0].shape == (4, ) + np.testing.assert_allclose(out[0], 0.5 * (grid[0][:-1] + grid[0][1:])) + + def test_1d_already_cell_centered_passthrough(self): + grid = [np.linspace(0.1, 0.9, 4)] # already 4 cell centers + out = nodal_to_cell_centered_grid(grid, cells=np.array([4])) + np.testing.assert_allclose(out[0], grid[0]) + + def test_2d_nodal_grids(self): + grid = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 2.0, 4)] + out = nodal_to_cell_centered_grid(grid, cells=np.array([4, 3])) + assert len(out) == 2 + assert out[0].shape == (4, ) + assert out[1].shape == (3, ) + + def test_dimension_mismatch_raises(self): + grid = [np.linspace(0.0, 1.0, 5)] + with pytest.raises(ValueError, match="doesn't match"): + nodal_to_cell_centered_grid(grid, cells=np.array([4, 3])) + + def test_bad_axis_length_raises(self): + grid = [np.linspace(0.0, 1.0, 6)] # neither 4 nor 5 points + with pytest.raises(ValueError, match="terribly wrong"): + nodal_to_cell_centered_grid(grid, cells=np.array([4])) + + def test_meshgrid_true_returns_ij_indexed_grid(self): + grid = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 2.0, 4)] + out = nodal_to_cell_centered_grid(grid, + cells=np.array([4, 3]), + meshgrid=True) + assert len(out) == 2 + assert out[0].shape == (4, 3) + assert out[1].shape == (4, 3) + + def test_meshgrid_false_keeps_1d_axes(self): + grid = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 2.0, 4)] + out = nodal_to_cell_centered_grid(grid, + cells=np.array([4, 3]), + meshgrid=False) + assert out[0].ndim == 1 + assert out[1].ndim == 1 + + def test_meshgrid_ignored_for_1d(self): + grid = [np.linspace(0.0, 1.0, 5)] + out = nodal_to_cell_centered_grid(grid, cells=np.array([4]), meshgrid=True) + assert len(out) == 1 + assert out[0].ndim == 1 + + def test_2d_array_grid_nodal(self): + """Multi-dimensional (already-meshgridded) coordinate arrays: the + 2-D-shaped-grid branch of the function.""" + x_nodal = np.linspace(0.0, 1.0, 5) + y_nodal = np.linspace(0.0, 2.0, 4) + X, Y = np.meshgrid(x_nodal, y_nodal, indexing="ij") + out = nodal_to_cell_centered_grid([X, Y], cells=np.array([4, 3])) + assert out[0].shape == (4, 3) + assert out[1].shape == (4, 3) + + def test_2d_array_grid_already_cell_centered_passthrough(self): + """Multi-dimensional grid array whose axis already matches ``cells`` + (no averaging needed) -- the ``grid[d].shape[d] == cells[d]`` + passthrough branch for array-shaped (already-meshgridded) grids.""" + x_cc = np.linspace(0.1, 0.9, 4) + y_cc = np.linspace(0.2, 1.8, 3) + X, Y = np.meshgrid(x_cc, y_cc, indexing="ij") + out = nodal_to_cell_centered_grid([X, Y], cells=np.array([4, 3])) + np.testing.assert_allclose(out[0], X) + np.testing.assert_allclose(out[1], Y) + + def test_multidim_grid_array_with_single_dimension(self): + """A multi-dimensional (ndim > 1) coordinate array in a 1-D grid (the + ``num_dims == 1`` branch of the array-shaped-grid case): averaging + happens along axis 0 only, leaving the other axis untouched.""" + grid = [ + np.array([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0], [8.0, 9.0]]) + ] + out = nodal_to_cell_centered_grid(grid, cells=np.array([4])) + assert len(out) == 1 + assert out[0].shape == (4, 2) + np.testing.assert_allclose(out[0], + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]) + + def test_2d_array_grid_bad_shape_raises(self): + x_nodal = np.linspace(0.0, 1.0, 6) # neither 4 nor 5 along axis 0 + y_nodal = np.linspace(0.0, 2.0, 4) + X, Y = np.meshgrid(x_nodal, y_nodal, indexing="ij") + with pytest.raises(ValueError, match="terribly wrong"): + nodal_to_cell_centered_grid([X, Y], cells=np.array([4, 3])) diff --git a/tests/test_numerics_misc.py b/tests/test_numerics_misc.py new file mode 100644 index 00000000..a219c01d --- /dev/null +++ b/tests/test_numerics_misc.py @@ -0,0 +1,130 @@ +"""Tests for postgkyl.numerics.mag_sq / rel_change / rotation_matrix.""" + +from __future__ import annotations + +import numpy as np + +from postgkyl.numerics.mag_sq import mag_sq +from postgkyl.numerics.rel_change import rel_change +from postgkyl.numerics.rotation_matrix import rotation_matrix + +_G1 = [np.array([0.0, 1.0])] + +# --------------------------------------------------------------------------- +# mag_sq +# --------------------------------------------------------------------------- + + +class TestMagSq: + + def test_unit_x_vector(self): + _, out = mag_sq(_G1, np.array([[1.0, 0.0, 0.0]])) + np.testing.assert_allclose(out.flat[0], 1.0) + + def test_3_4_0_vector(self): + _, out = mag_sq(_G1, np.array([[3.0, 4.0, 0.0]])) + np.testing.assert_allclose(out.flat[0], 25.0) + + def test_output_has_trailing_dim(self): + _, out = mag_sq(_G1, np.array([[1.0, 2.0, 3.0]])) + assert out.ndim == 2 + assert out.shape[-1] == 1 + + def test_custom_coords(self): + _, out = mag_sq(_G1, + np.array([[0.0, 0.0, 0.0, 3.0, 4.0, 0.0]]), + coords="3:6") + np.testing.assert_allclose(out.flat[0], 25.0) + + def test_multi_cell(self): + grid = [np.linspace(0.0, 1.0, 4)] + values = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]) + _, out = mag_sq(grid, values) + np.testing.assert_allclose(out[:, 0], [1.0, 1.0, 2.0]) + + def test_grid_returned_unchanged(self): + grid = [np.linspace(0.0, 1.0, 3)] + out_grid, _ = mag_sq(grid, np.array([[1.0, 0.0, 0.0]])) + np.testing.assert_allclose(out_grid[0], grid[0]) + + +# --------------------------------------------------------------------------- +# rel_change +# --------------------------------------------------------------------------- + + +class TestRelChange: + + def test_doubled_values(self): + grid = [np.linspace(0.0, 1.0, 4)] + v0 = np.array([[1.0], [2.0], [3.0]]) + v1 = np.array([[2.0], [4.0], [6.0]]) + _, out = rel_change(grid, v0, v1) + np.testing.assert_allclose(out[:, 0], [1.0, 1.0, 1.0]) + + def test_no_change_gives_zero(self): + grid = [np.linspace(0.0, 1.0, 4)] + v = np.array([[1.0], [2.0], [3.0]]) + _, out = rel_change(grid, v.copy(), v.copy()) + np.testing.assert_allclose(out[:, 0], 0.0, atol=1e-14) + + def test_with_comp_normalizes_by_selected_component(self): + grid = [np.linspace(0.0, 1.0, 3)] + v0 = np.array([[2.0, 4.0], [1.0, 2.0]]) + v1 = np.array([[4.0, 8.0], [2.0, 4.0]]) + _, out = rel_change(grid, v0, v1, comp=0) + np.testing.assert_allclose(out[0, 0], 1.0) + np.testing.assert_allclose(out[0, 1], 2.0) + + def test_multi_component(self): + grid = [np.linspace(0.0, 1.0, 3)] + v0 = np.array([[1.0, 2.0], [1.0, 4.0]]) + v1 = np.array([[2.0, 4.0], [3.0, 8.0]]) + _, out = rel_change(grid, v0, v1) + np.testing.assert_allclose(out[0, 0], 1.0) + np.testing.assert_allclose(out[0, 1], 1.0) + np.testing.assert_allclose(out[1, 0], 2.0) + np.testing.assert_allclose(out[1, 1], 1.0) + + +# --------------------------------------------------------------------------- +# rotation_matrix +# --------------------------------------------------------------------------- + + +class TestRotationMatrix: + + def test_basic_shape(self): + v = np.array([1.0, 2.0, 3.0]) + R = rotation_matrix(v) + assert R.shape == (3, 3) + + def test_returns_ndarray(self): + v = np.array([1.0, 2.0, 3.0]) + assert isinstance(rotation_matrix(v), np.ndarray) + + def test_arbitrary_vector_first_row_is_direction(self): + v = np.array([3.0, 4.0, 1.0]) + R = rotation_matrix(v) + k = v / np.abs(v) + np.testing.assert_allclose(R[0], k, atol=1e-10) + + def test_positive_vector(self): + v = np.array([1.0, 2.0, 3.0]) + R = rotation_matrix(v) + np.testing.assert_allclose(R[0], np.array([1.0, 1.0, 1.0]), atol=1e-10) + + def test_returns_non_zero_matrix(self): + v = np.array([1.0, 2.0, 3.0]) + R = rotation_matrix(v) + assert R.dtype == float + assert np.any(R != 0) + + def test_rows_are_mutually_orthogonal(self): + """Analytic check: rotation_matrix builds an (unnormalized) orthogonal + frame -- each row should be perpendicular to every other row.""" + v = np.array([2.0, -3.0, 5.0]) + R = rotation_matrix(v) + np.testing.assert_allclose(R[0] @ R[1], 0.0, atol=1e-10) + np.testing.assert_allclose(R[0] @ R[2], 0.0, atol=1e-10) + np.testing.assert_allclose(R[1] @ R[2], 0.0, atol=1e-10) diff --git a/tests/test_operations_animate.py b/tests/test_operations_animate.py new file mode 100644 index 00000000..5a8d9645 --- /dev/null +++ b/tests/test_operations_animate.py @@ -0,0 +1,77 @@ +"""Tests for the canonical ``render.animate`` callable through its exact +``operations.animate`` alias (mirrors the aliases of ``render.plot``).""" + +from __future__ import annotations + +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") +pytestmark = needs_gkeyll + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GEN = os.path.join(DATA, "generated") +F1D = os.path.join(GEN, "1d_ms_p1.gkyl") + + +@pytest.fixture(autouse=True) +def _close_figs(): + plt.close("all") + yield + plt.close("all") + + +def _three_interpolated_frames(): + return [pg.load(F1D).interpolate().select(comp=c) for c in (0, 0, 0)] + + +class TestAnimateVerb: + + def test_already_interpolated_frames_pass_through(self): + from matplotlib.animation import FuncAnimation + anim = operations.animate(_three_interpolated_frames(), no_show=True) + assert isinstance(anim, FuncAnimation) + assert anim._save_count == 3 + + def test_modal_frames_are_materialized_first(self): + """A raw (non-interpolated) modal dataset is bridged through its NumPy + shadow (nodal value_form), just like ``render.plot``.""" + from matplotlib.animation import FuncAnimation + a = pg.load(F1D).to_nodal() + b = pg.load(F1D).to_nodal() + anim = operations.animate([a, b], no_show=True) + assert isinstance(anim, FuncAnimation) + assert anim._save_count == 2 + + def test_raw_modal_frame_without_representation_raises(self): + a = pg.load(F1D) # still modal coefficients + with pytest.raises(ValueError, match="not plottable"): + operations.animate([a], no_show=True) + + def test_grouped_frames_preserve_structure(self): + from matplotlib.animation import FuncAnimation + a = pg.load(F1D).interpolate() + b = pg.load(F1D).interpolate() + c = pg.load(F1D).interpolate() + anim = operations.animate([[a, b], [c]], no_show=True) + assert isinstance(anim, FuncAnimation) + assert anim._save_count == 2 + + def test_saveframes_end_to_end(self, tmp_path): + prefix = str(tmp_path / "frame") + paths = operations.animate(_three_interpolated_frames(), + saveframes=prefix, + no_show=True) + assert len(paths) == 3 + for p in paths: + assert os.path.isfile(p) diff --git a/tests/test_operations_collect.py b/tests/test_operations_collect.py new file mode 100644 index 00000000..912b62e1 --- /dev/null +++ b/tests/test_operations_collect.py @@ -0,0 +1,110 @@ +"""Tests for the ``collect`` verb -- stacking many datasets onto a time axis.""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _frame(time, value, grid=None): + grid = grid if grid is not None else [np.linspace(0.0, 1.0, 5)] + d = GDataState(ctx={"time": time}) + d.push(list(grid), np.full((4, 1), value)) + return d + + +def test_stacks_frames_sorted_by_time(): + a = _frame(1.0, 2.0) + b = _frame(0.0, 1.0) + out = operations.collect(a, b) + np.testing.assert_allclose(out.get_grid()[0], [0.0, 1.0]) + np.testing.assert_allclose(out.get_values()[0].flatten(), 1.0) + np.testing.assert_allclose(out.get_values()[1].flatten(), 2.0) + + +def test_accepts_a_list_argument(): + frames = [_frame(0.0, 1.0), _frame(1.0, 2.0)] + out = operations.collect(frames) + assert out.get_values().shape[0] == 2 + + +def test_sumdata_reduces_spatial_axes(): + a = _frame(0.0, 3.0) + b = _frame(1.0, 5.0) + out = operations.collect(a, b, sumdata=True) + np.testing.assert_allclose(out.get_values().flatten(), [3.0 * 4, 5.0 * 4]) + assert out.get_grid()[0].shape == (2, ) + + +def test_frame_stamp_falls_back_to_position_when_no_time_or_frame(): + a = GDataState() + a.push([np.linspace(0.0, 1.0, 5)], np.full((4, 1), 10.0)) + b = GDataState() + b.push([np.linspace(0.0, 1.0, 5)], np.full((4, 1), 20.0)) + out = operations.collect(a, b) + np.testing.assert_allclose(out.get_grid()[0], [0, 1]) + + +def test_period_folds_time_axis(): + a = _frame(0.0, 1.0) + b = _frame(3.0, 2.0) # 3.0 % 2.0 == 1.0 + out = operations.collect(a, b, period=2.0) + np.testing.assert_allclose(sorted(out.get_grid()[0]), [0.0, 1.0]) + + +def test_tag_and_label_defaults(): + a, b = _frame(0.0, 1.0), _frame(1.0, 2.0) + out = operations.collect(a, b) + assert out.get_tag() == "default" + assert out.get_label() == "collect" + + +def test_tag_and_label_explicit(): + a, b = _frame(0.0, 1.0), _frame(1.0, 2.0) + out = operations.collect(a, b, tag="series", label="my series") + assert out.get_tag() == "series" + assert out.get_label() == "my series" + + +def test_empty_raises(): + with pytest.raises(ValueError): + operations.collect() + + +def test_chunk_splits_into_multiple_datasets(): + frames = [_frame(float(i), float(i)) for i in range(5)] + out = operations.collect(frames, chunk=2) + assert isinstance(out, list) + assert len(out) == 3 + np.testing.assert_allclose(out[0].get_grid()[0], [0.0, 1.0]) + np.testing.assert_allclose(out[1].get_grid()[0], [2.0, 3.0]) + np.testing.assert_allclose(out[2].get_grid()[0], [4.0]) + + +def test_chunk_falsy_returns_single_dataset(): + frames = [_frame(0.0, 1.0), _frame(1.0, 2.0)] + out = operations.collect(frames, chunk=0) + assert not isinstance(out, list) + assert out.get_values().shape[0] == 2 + + +@needs_gkeyll +def test_rejects_modal_data(): + modal = pg.load(F1) + numpy_side = _frame(0.0, 1.0) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.collect(modal, numpy_side) diff --git a/tests/test_operations_curvilinear_diff_integrate.py b/tests/test_operations_curvilinear_diff_integrate.py new file mode 100644 index 00000000..d822aaa2 --- /dev/null +++ b/tests/test_operations_curvilinear_diff_integrate.py @@ -0,0 +1,170 @@ +"""``differentiate``/``integrate`` on curvilinear (``.map(space="conf")``) +grids -- both verbs need the block's Jacobian (``numerics.curvilinear``) +instead of treating each axis as separable. Uses exactly linear coordinate +maps (rotation, shear) so the finite-difference machinery reproduces the +analytic answer up to floating-point precision, independent of resolution. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.numerics import curvilinear + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") +pytestmark = needs_gkeyll + + +# --------------------------------------------------------------- test helpers +def _project_2d(fn, lower, upper, cells, basis_type, poly_order): + """Exact per-cell modal coefficients of ``fn(z0, z1)`` (mirrors + ``tests/test_operations_map.py``'s helper of the same name).""" + node_eta = gpython.basis.node_coords(basis_type, 2, poly_order) + n2m = gpython.basis.nodal_to_modal_matrix(basis_type, 2, poly_order) + dz = [(upper[d] - lower[d]) / cells[d] for d in range(2)] + c0 = lower[0] + (np.arange(cells[0]) + 0.5) * dz[0] + c1 = lower[1] + (np.arange(cells[1]) + 0.5) * dz[1] + centers = np.stack(np.meshgrid(c0, c1, indexing="ij"), axis=-1) + node_phys = ( + centers[:, :, None, :] + + 0.5 * np.array(dz)[None, None, None, :] * node_eta[None, None, :, :]) + nodal_vals = fn(node_phys[..., 0], node_phys[..., 1]) + return np.einsum("ij,...j->...i", n2m, nodal_vals) + + +def _synthetic_map(coeffs, + lower, + upper, + cells, + *, + basis_type="serendipity", + poly_order=1): + d = GDataState() + d.ctx.update(basis_type=basis_type, + poly_order=poly_order, + value_form="modal", + cells=np.asarray(cells, dtype=np.int64)) + grid = [ + np.linspace(lower[i], upper[i], + int(cells[i]) + 1) for i in range(len(cells)) + ] + d.push(grid, gpython.GkylArray.from_numpy(coeffs)) + return d + + +def _linear_map_dataset(a00, + a01, + a10, + a11, + *, + lower=(0.0, 0.0), + upper=(4.0, 4.0), + cells=(8, 8), + poly_order=1): + """A target dataset whose grid is deformed by the linear map + ``x = a00*z0 + a01*z1``, ``y = a10*z0 + a11*z1`` -- exact for any + ``poly_order >= 1`` (an affine function).""" + m0 = _project_2d(lambda z0, z1: a00 * z0 + a01 * z1, lower, upper, cells, + "serendipity", poly_order) + m1 = _project_2d(lambda z0, z1: a10 * z0 + a11 * z1, lower, upper, cells, + "serendipity", poly_order) + modal = np.concatenate([m0, m1], axis=-1) + mapping = _synthetic_map(modal, lower, upper, cells, poly_order=poly_order) + target = GDataState() + target.push([np.linspace(lower[i], upper[i], cells[i] + 1) for i in range(2)], + np.zeros(tuple(cells) + (1, ))) + mapped = operations.map(target, mapping, space="conf") + + x = curvilinear.cell_center(mapped.grid[0]) + y = curvilinear.cell_center(mapped.grid[1]) + return mapped, x, y + + +# --------------------------------------------------------------- differentiate +class TestDifferentiateCurvilinear: + + def test_rotation_gradient_matches_analytic(self): + """A rigid rotation: physical gradient of f = x**2 + y is exactly + (2x, 1) regardless of the rotation angle.""" + theta = 0.3 + c, s = np.cos(theta), np.sin(theta) + mapped, x, y = _linear_map_dataset(c, -s, s, c) + field = mapped._result(mapped.grid, (x**2 + y)[..., np.newaxis]) + + dfdx = operations.differentiate(field, direction=0) + dfdy = operations.differentiate(field, direction=1) + np.testing.assert_allclose(dfdx.values[..., 0], 2 * x, atol=1e-8) + np.testing.assert_allclose(dfdy.values[..., 0], np.ones_like(y), atol=1e-8) + + def test_shear_gradient_of_linear_field_is_exact(self): + """A non-orthogonal (shear) map: for a field that is itself linear in + the physical coordinates, the chain-rule reconstruction is exact.""" + mapped, x, y = _linear_map_dataset(1.0, 0.5, 0.0, 1.0) + a, b = 3.0, -2.0 + field = mapped._result(mapped.grid, (a * x + b * y)[..., np.newaxis]) + + grad = operations.differentiate(field) + np.testing.assert_allclose(grad.values[..., 0], a, atol=1e-8) + np.testing.assert_allclose(grad.values[..., 1], b, atol=1e-8) + + def test_full_gradient_direction_none_matches_per_direction_calls(self): + mapped, x, y = _linear_map_dataset(1.0, 0.3, -0.2, 1.0) + field = mapped._result(mapped.grid, (x**2 + y**2)[..., np.newaxis]) + + full = operations.differentiate(field) + dx = operations.differentiate(field, direction=0) + dy = operations.differentiate(field, direction=1) + np.testing.assert_allclose(full.values[..., 0], dx.values[..., 0]) + np.testing.assert_allclose(full.values[..., 1], dy.values[..., 0]) + + +# ------------------------------------------------------------------ integrate +class TestIntegrateCurvilinear: + + def test_rotation_preserves_area(self): + """Integrating the constant field 1 over a rotated square recovers the + original (unrotated) domain area -- rotation preserves area.""" + theta = 0.4 + c, s = np.cos(theta), np.sin(theta) + lower, upper, cells = (0.0, 0.0), (4.0, 2.0), (16, 16) + mapped, x, _y = _linear_map_dataset(c, + -s, + s, + c, + lower=lower, + upper=upper, + cells=cells) + field = mapped._result(mapped.grid, np.ones(tuple(cells) + (1, ))) + + total = operations.integrate(field, "0,1") + expected_area = (upper[0] - lower[0]) * (upper[1] - lower[1]) + np.testing.assert_allclose(total, expected_area, rtol=1e-6) + + def test_shear_area_matches_determinant(self): + """A shear map's area scales by the (constant) Jacobian determinant of + the linear transform, |a00*a11 - a01*a10|.""" + a00, a01, a10, a11 = 1.0, 0.5, 0.0, 1.0 + lower, upper, cells = (0.0, 0.0), (2.0, 3.0), (10, 12) + mapped, x, _y = _linear_map_dataset(a00, + a01, + a10, + a11, + lower=lower, + upper=upper, + cells=cells) + field = mapped._result(mapped.grid, np.ones(tuple(cells) + (1, ))) + + total = operations.integrate(field, "0,1") + det = abs(a00 * a11 - a01 * a10) + expected_area = det * (upper[0] - lower[0]) * (upper[1] - lower[1]) + np.testing.assert_allclose(total, expected_area, rtol=1e-6) + + def test_partial_block_reduction_raises(self): + mapped, x, y = _linear_map_dataset(1.0, 0.3, -0.2, 1.0) + field = mapped._result(mapped.grid, (x + y)[..., np.newaxis]) + with pytest.raises(ValueError, match="curvilinear"): + operations.integrate(field, "0") diff --git a/tests/test_operations_differentiate.py b/tests/test_operations_differentiate.py new file mode 100644 index 00000000..6f0d47b1 --- /dev/null +++ b/tests/test_operations_differentiate.py @@ -0,0 +1,105 @@ +"""Tests for the ``differentiate`` verb -- numerical gradient of field data. + +Per the layer-03 differentiate-decision note, this is a post-``.interpolate()`` +verb: it takes NumPy field values and refuses native modal (gkyl-backed) +data, exactly like ``select``. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +def _quadratic_1d(n=40): + edges = np.linspace(0.0, 1.0, n + 1) + centers = 0.5 * (edges[:-1] + edges[1:]) + y = centers**2 # d/dx = 2x + return _make([edges], y[:, np.newaxis]), centers + + +def test_full_gradient_matches_analytic_derivative_1d(): + d, centers = _quadratic_1d() + out = operations.differentiate(d) + np.testing.assert_allclose(out.get_values().flatten(), + 2.0 * centers, + atol=1e-2) + assert out.get_num_comps() == 1 # 1 comp * 1 dim = 1 + + +def test_direction_matches_full_gradient_in_1d(): + d, _ = _quadratic_1d() + full = operations.differentiate(d) + by_dir = operations.differentiate(d, direction=0) + np.testing.assert_allclose(full.get_values(), by_dir.get_values()) + + +def test_grid_unchanged(): + d, _ = _quadratic_1d() + out = operations.differentiate(d) + np.testing.assert_allclose(out.get_grid()[0], d.get_grid()[0]) + + +def test_2d_full_gradient_stacks_components(): + e0 = np.linspace(0.0, 1.0, 21) + e1 = np.linspace(0.0, 1.0, 21) + c0 = 0.5 * (e0[:-1] + e0[1:]) + c1 = 0.5 * (e1[:-1] + e1[1:]) + X, Y = np.meshgrid(c0, c1, indexing="ij") + values = (X**2 + Y)[..., np.newaxis] # d/dx = 2x, d/dy = 1 + d = _make([e0, e1], values) + out = operations.differentiate(d) + assert out.get_num_comps() == 2 + np.testing.assert_allclose(out.get_values()[..., 0], 2 * X, atol=1e-2) + np.testing.assert_allclose(out.get_values()[..., 1], + np.ones_like(Y), + atol=1e-2) + + single = operations.differentiate(d, direction=1) + np.testing.assert_allclose(single.get_values()[..., 0], + np.ones_like(Y), + atol=1e-2) + + +def test_inplace_and_tag_label(): + d, _ = _quadratic_1d() + out = operations.differentiate(d, tag="grad", label="dq/dx", inplace=True) + assert out is d + assert d.get_tag() == "grad" + assert d.get_label() == "dq/dx" + + +def test_mismatched_grid_length_raises(): + # Cell-centered grid (matches value count, not the expected nodal edges) + # cannot form the required cell-widths -- this is the documented caveat. + x = np.linspace(0.0, 1.0, 10) + d = _make([x], (x**2)[:, np.newaxis]) + with pytest.raises(ValueError): + operations.differentiate(d) + + +@needs_gkeyll +def test_rejects_modal_data(): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.differentiate(d) diff --git a/tests/test_operations_evaluate.py b/tests/test_operations_evaluate.py new file mode 100644 index 00000000..0cf31b0a --- /dev/null +++ b/tests/test_operations_evaluate.py @@ -0,0 +1,296 @@ +"""Tests for the ``evaluate`` verb -- the RPN expression evaluator over datasets.""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +def _field(value, grid=None): + grid = grid if grid is not None else [np.linspace(0.0, 1.0, 5)] + return _make(grid, np.full((4, 1), value)) + + +# -------------------------------------------------------- parity with verbs +def test_add_two_datasets_matches_direct_arithmetic(): + """The grammar's dataset-index tokens are plain 'fN' (no brackets -- + 'fN[c]' is the *component* selector, per the module docstring); this is + the byte-compatible spelling for combining two whole datasets.""" + a, b = _field(2.0), _field(3.0) + out = operations.evaluate("f0 f1 +", a, b) + np.testing.assert_allclose(out.get_values().flatten(), 5.0) + + +def test_default_f_means_f0(): + a = _field(4.0) + out = operations.evaluate("f 2 *", a) + np.testing.assert_allclose(out.get_values().flatten(), 8.0) + + +def test_component_bracket_selects_a_component(): + a = _make([np.linspace(0.0, 1.0, 5)], np.tile([1.0, 2.0, 3.0], (4, 1))) + out = operations.evaluate("f0[1] sq", a) + np.testing.assert_allclose(out.get_values().flatten(), 4.0) + + +def test_ctx_key_token(): + a = _field(1.0) + a.ctx["scale"] = 3.0 + out = operations.evaluate("f0 f0.scale *", a) + np.testing.assert_allclose(out.get_values().flatten(), 3.0) + + +def test_unknown_ctx_key_raises(): + a = _field(1.0) + with pytest.raises(ValueError, match="unknown ctx key"): + operations.evaluate("f0.nope", a) + + +# -------------------------------------------------------------- operators +def test_sqrt_and_abs(): + a = _field(-4.0) + out = operations.evaluate("f abs sqrt", a) + np.testing.assert_allclose(out.get_values().flatten(), 2.0) + + +def test_min_max_mean(): + a = _make([np.linspace(0.0, 1.0, 5)], + np.array([1.0, 2.0, 3.0, 4.0])[:, np.newaxis]) + assert operations.evaluate("f min", + a).get_values().flatten()[0] == pytest.approx(1.0) + assert operations.evaluate("f max", + a).get_values().flatten()[0] == pytest.approx(4.0) + assert operations.evaluate("f mean", + a).get_values().flatten()[0] == pytest.approx(2.5) + + +def test_numeric_literal_and_axis_slice_literal(): + a = _field(2.0) + out = operations.evaluate("f 3.0 +", a) + np.testing.assert_allclose(out.get_values().flatten(), 5.0) + + +# ------------------------------------------------------------------ result +def test_result_class_and_defaults(): + a, b = _field(2.0), _field(3.0) + out = operations.evaluate("f0 f1 +", a, b) + assert isinstance(out, GDataState) + assert out.get_tag() == "default" + assert out.get_label() == "f0 f1 +" + + +def test_tag_and_label_explicit(): + a, b = _field(2.0), _field(3.0) + out = operations.evaluate("f0 f1 +", a, b, tag="t", label="sum") + assert out.get_tag() == "t" + assert out.get_label() == "sum" + + +def test_num_comps_reflects_the_actual_output_not_a_stale_operand_value(): + """A component-changing op (here 'dot', which reduces a vector to a + scalar) must not have its output metadata clobbered by a stale + 'num_comps'/'cells' merged in from the (differently-shaped) operands.""" + a = _make([np.linspace(0.0, 1.0, 5)], np.tile([1.0, 0.0, 0.0], (4, 1))) + b = _make([np.linspace(0.0, 1.0, 5)], np.tile([1.0, 0.0, 0.0], (4, 1))) + out = operations.evaluate("f0 f1 dot", a, b) + assert out.get_num_comps() == 1 + np.testing.assert_allclose(out.get_values().flatten(), 1.0) + + +def test_conflicting_ctx_keys_are_dropped_not_merged(): + a = _field(2.0) + b = _field(3.0) + a.ctx["note"] = "A" + b.ctx["note"] = "B" + out = operations.evaluate("f0 f1 +", a, b) + assert "note" not in out.ctx + + +def test_bracket_literal_and_colon_axis_literal(): + a = _make([np.linspace(0.0, 1.0, 5)], + np.array([1.0, 2.0, 3.0, 4.0])[:, np.newaxis]) + # a bare bracket literal (no leading 'f') exercises the eval() fallback + out = operations.evaluate("[1,2,3] mean", a) + np.testing.assert_allclose(out.get_values().flatten(), 2.0) + # a bare colon axis spec exercises the str-literal fallback + 'int' + out2 = operations.evaluate("f 0:1 int", a) + assert isinstance(out2, GDataState) + + +# -------------------------------------------------------------------- errors +def test_empty_datasets_raises(): + with pytest.raises(ValueError, match="at least one dataset"): + operations.evaluate("f 2 *") + + +def test_empty_expression_raises(): + a = _field(1.0) + with pytest.raises(ValueError, match="produced no result"): + operations.evaluate("", a) + + +def test_unrecognized_token_raises(): + a = _field(1.0) + with pytest.raises(ValueError, match="neither data nor an operator"): + operations.evaluate("f totally_bogus_token", a) + + +def test_operator_failure_is_wrapped_in_value_error(): + # 1D grid (num_dims=1) with 4 components: 'div' (num_in=1) refuses a + # component count larger than the number of dimensions. + a = _make([np.linspace(0.0, 1.0, 2)], np.tile([1.0, 2.0, 3.0, 4.0], (1, 1))) + with pytest.raises(ValueError, match="ERROR in 'evaluate div'"): + operations.evaluate("f div", a) + + +@needs_gkeyll +def test_modal_data_supported_ops_use_weak_kernels(): + """+ - * / and integer pow/sq have an exact DG meaning, so they run on raw + modal coefficients via Gkeyll's own weak kernels -- the result stays + native (gkyl-backed), never silently dropping to plain NumPy math.""" + d = pg.load(F1) + assert d.backend == "gkyl" + for chain in ("f sq", "f 2 *", "1 f /", "f0 f0 +"): + result = operations.evaluate(chain, d) + assert result.backend == "gkyl" + + +@needs_gkeyll +def test_modal_data_unsupported_op_warns_and_falls_back(): + """sqrt has no weak-kernel form: rather than hard-blocking (basis/ + value_form metadata can be wrong), evaluate warns and computes on the + raw coefficient view, which is exact only when coefficient 0 already IS + the point value.""" + d = pg.load(F1) + with pytest.warns(UserWarning, match="weak-kernel"): + result = operations.evaluate("f sqrt", d) + assert result.backend == "numpy" + + +def test_available_operator_vocabulary_is_sorted_and_public(): + operators = pg.available_evaluate_operators() + assert operators == sorted(operators) + assert {"+", "pow", "sqrt", "int"} <= set(operators) + + +@needs_gkeyll +def test_modal_dataset_binary_operators_cover_each_weak_kernel_dispatch(): + d = _native_field(2.0, "modal") + for token in ("-", "*", "/"): + result = operations.evaluate(f"f0 f0 {token}", d, d) + assert result.backend == "gkyl" + + with pytest.warns(UserWarning, match="between two modal datasets"): + result = operations.evaluate("f0 f0 pow", d, d) + assert result.backend == "numpy" + + with pytest.warns(UserWarning, match="no weak-kernel form"): + result = operations.evaluate("f0 f0 max2", d, d) + assert result.backend == "numpy" + + +@needs_gkeyll +def test_modal_dataset_scalar_operators_cover_operand_order_and_power(): + d = _native_field(2.0, "modal") + for chain in ("f 2 +", "f 2 -", "2 f -", "f 2 pow"): + result = operations.evaluate(chain, d) + assert result.backend == "gkyl" + + for chain in ("f 0 pow", "2 f pow"): + with pytest.warns(UserWarning, match="positive integer"): + result = operations.evaluate(chain, d) + assert result.backend == "numpy" + + +@needs_gkeyll +def test_modal_dispatch_warns_for_missing_or_mismatched_metadata(): + d = _native_field(2.0, "modal") + missing = d.clone() + missing.ctx.pop("basis_type") + with pytest.warns(UserWarning, match="no basis_type/poly_order"): + operations.evaluate("f sq", missing) + + mismatched = d.clone() + mismatched.ctx["basis_type"] = "tensor" + with pytest.warns(UserWarning, match="different DG bases"): + operations.evaluate("f0 f1 +", d, mismatched) + + with pytest.warns(UserWarning, match="plain array"): + operations.evaluate("f [1] +", d) + + +@needs_gkeyll +def test_modal_dispatch_rejects_operators_without_a_matching_arity(): + d = _native_field(2.0, "modal") + with pytest.warns(UserWarning, match="3 operands"): + result = operations.evaluate("f 0 2 scale_comp", d) + assert result.backend == "numpy" + + +def _native_field(value, value_form): + grid = [np.linspace(0.0, 1.0, 5)] + native = gpython.GkylArray.from_numpy(np.full((4, 1), value)) + return _make(grid, + native, + basis_type="serendipity", + poly_order=0, + value_form=value_form) + + +@needs_gkeyll +@pytest.mark.parametrize("value_form", ["nodal", "quad"]) +def test_native_pointwise_evaluation_stays_native(value_form): + d = _native_field(4.0, value_form) + result = operations.evaluate("f sqrt", d) + assert result.backend == "gkyl" + assert result.ctx["value_form"] == value_form + np.testing.assert_allclose(result.values, 2.0) + + +@needs_gkeyll +def test_native_point_reduction_leaves_the_value_form_domain(): + d = _native_field(4.0, "nodal") + result = operations.evaluate("f mean", d) + assert result.backend == "numpy" + assert "value_form" not in result.ctx + assert result.ctx["interpolated"] is True + + +@needs_gkeyll +def test_mixed_native_point_value_forms_warn_and_fall_back(): + nodal = _native_field(2.0, "nodal") + quad = _native_field(3.0, "quad") + with pytest.warns(UserWarning, match="different value_forms"): + result = operations.evaluate("f0 f1 +", nodal, quad) + assert result.backend == "numpy" + np.testing.assert_allclose(result.values, 5.0) + + +@needs_gkeyll +def test_modal_dispatch_scalar_helpers_cover_scalar_shapes(): + from importlib import import_module + evaluate_module = import_module("postgkyl.operations.evaluate") + assert evaluate_module._as_scalar(np.int64(3)) == 3.0 + assert evaluate_module._as_scalar(np.array([3.0])) is None + assert evaluate_module._modal_kernel("+", [None], [np.array([1.0])], + [{}]) is None diff --git a/tests/test_operations_field.py b/tests/test_operations_field.py new file mode 100644 index 00000000..aa8fd368 --- /dev/null +++ b/tests/test_operations_field.py @@ -0,0 +1,310 @@ +"""Tests for the small field-domain operations verbs: fft, magsq, relchange, mask, +grid, val2coord, extract_input. +""" + +from __future__ import annotations + +import base64 +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastategroup import GDataStateGroup +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +# ============================================================== operations.fft +class TestFft: + + def test_analytic_sine_peak(self): + N = 32 + edges = np.linspace(0.0, 1.0, N + 1) + x_cc = 0.5 * (edges[:-1] + edges[1:]) + f0 = 4 + values = np.sin(2 * np.pi * f0 * x_cc)[:, np.newaxis] + d = _make([edges], values) + out = operations.fft(d) + assert isinstance(out, GDataState) + freq = out.get_grid()[0] + ft = out.get_values() + peak = freq[np.argmax(np.abs(ft[:, 0]))] + assert abs(abs(peak) - f0) < 1e-9 + + def test_psd_returns_positive_frequencies_only(self): + N = 16 + d = _make([np.linspace(0.0, 1.0, N + 1)], np.ones((N, 1))) + out = operations.fft(d, psd=True) + assert out.get_values().shape[0] == N // 2 + + def test_inplace_mutates(self): + d = _make([np.linspace(0.0, 1.0, 17)], np.ones((16, 1))) + out = operations.fft(d, inplace=True) + assert out is d + + def test_tag_and_label(self): + d = _make([np.linspace(0.0, 1.0, 17)], np.ones((16, 1))) + out = operations.fft(d, tag="spec", label="lbl") + assert out.get_tag() == "spec" + assert out.get_label() == "lbl" + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.fft(d) + + +# ============================================================ operations.magsq +class TestMagsq: + + def _vec3(self): + return _make([np.linspace(0.0, 1.0, 5)], np.tile([1.0, 2.0, 3.0], (4, 1))) + + def test_value_and_num_comps(self): + out = operations.magsq(self._vec3()) + np.testing.assert_allclose(out.get_values().flat[0], 14.0) # 1+4+9 + assert out.get_num_comps() == 1 + + def test_custom_coords(self): + out = operations.magsq(self._vec3(), coords="1:3") + np.testing.assert_allclose(out.get_values().flat[0], 13.0) # 4+9 + + def test_inplace(self): + d = self._vec3() + assert operations.magsq(d, inplace=True) is d + + def test_tag(self): + out = operations.magsq(self._vec3(), tag="m") + assert out.get_tag() == "m" + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.magsq(d) + + +# ========================================================= operations.relchange +class TestRelchange: + + def test_value_componentwise(self): + grid = [np.linspace(0.0, 1.0, 5)] + ref = _make(grid, np.full((4, 1), 2.0)) + cur = _make(grid, np.full((4, 1), 3.0)) + out = operations.relchange(ref, cur) + np.testing.assert_allclose(out.get_values(), 0.5) # (3-2)/2 + + def test_value_with_explicit_comp(self): + grid = [np.linspace(0.0, 1.0, 5)] + ref = _make(grid, np.tile([2.0, 10.0], (4, 1))) + cur = _make(grid, np.tile([4.0, 4.0], (4, 1))) + out = operations.relchange(ref, cur, + comp=0) # normalize both by ref comp 0 (=2) + np.testing.assert_allclose(out.get_values()[..., 0], 1.0) # (4-2)/2 + np.testing.assert_allclose(out.get_values()[..., 1], -3.0) # (4-10)/2 + + def test_result_built_from_data_not_reference(self): + grid = [np.linspace(0.0, 1.0, 5)] + ref = _make(grid, np.full((4, 1), 2.0), tag="ref") + cur = _make(grid, np.full((4, 1), 3.0), tag="cur") + out = operations.relchange(ref, cur, tag="rc") + assert out.get_tag() == "rc" + + def test_inplace_mutates_data(self): + grid = [np.linspace(0.0, 1.0, 5)] + ref = _make(grid, np.full((4, 1), 2.0)) + cur = _make(grid, np.full((4, 1), 4.0)) + out = operations.relchange(ref, cur, inplace=True) + assert out is cur + + @needs_gkeyll + def test_rejects_modal_data(self): + grid = [np.linspace(0.0, 1.0, 5)] + numpy_side = _make(grid, np.full((4, 1), 2.0)) + modal = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.relchange(modal, numpy_side) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.relchange(numpy_side, modal) + + +# ============================================================== operations.mask +class TestMask: + + def _data(self): + return _make([np.linspace(0.0, 1.0, 6)], np.arange(5.0)[:, np.newaxis]) + + def test_mask_lower(self): + out = operations.mask(self._data(), lower=2.0) + assert np.ma.is_masked(out.get_values()) + assert out.get_values().mask[0, 0] + assert not out.get_values().mask[-1, 0] + + def test_mask_upper(self): + out = operations.mask(self._data(), upper=2.0) + assert out.get_values().mask[-1, 0] + + def test_mask_outside(self): + out = operations.mask(self._data(), lower=1.0, upper=3.0) + assert np.ma.is_masked(out.get_values()) + + def test_mask_from_dataset(self): + grid = [np.linspace(0.0, 1.0, 6)] + d = _make(grid, np.ones((5, 2))) + mask_field = _make(grid, np.array([[1.0], [-1.0], [1.0], [-1.0], [1.0]])) + out = operations.mask(d, mask_field) + values = out.get_values() + assert np.ma.is_masked(values) + assert values.mask[1, 0] and values.mask[1, 1] + assert not values.mask[0, 0] + + def test_mask_no_args_raises(self): + with pytest.raises(ValueError): + operations.mask(self._data()) + + def test_mask_from_dataset_multi_component_raises(self): + """mask_data must have exactly one component (see mask.py's docstring); + a multi-component mask does not "evenly divide" -- np.repeat produces + k*num_comps entries, which np.ma.masked_where rejects outright.""" + grid = [np.linspace(0.0, 1.0, 6)] + d = _make(grid, np.ones((5, 2))) + mask_field = _make( + grid, + np.array([[1.0, 1.0], [-1.0, -1.0], [1.0, 1.0], [-1.0, -1.0], + [1.0, 1.0]])) + with pytest.raises(IndexError): + operations.mask(d, mask_field) + + def test_inplace(self): + d = self._data() + out = operations.mask(d, lower=2.0, inplace=True) + assert out is d + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.mask(d, lower=0.0) + + +# ============================================================== operations.grid +class TestGrid: + + def test_1d_values_equal_grid(self): + edges = np.linspace(0.0, 1.0, 5) + d = _make([edges], np.ones((4, 1))) + out = operations.grid(d) + np.testing.assert_allclose(out.get_values()[..., 0], edges) + + def test_2d_meshgrid_shape(self): + edges = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 2.0, 4)] + d = _make(edges, np.ones((4, 3, 1))) + out = operations.grid(d) + assert out.get_num_comps() == 2 + assert out.get_values().shape == (5, 4, 2) + + def test_inplace(self): + edges = np.linspace(0.0, 1.0, 5) + d = _make([edges], np.ones((4, 1))) + out = operations.grid(d, inplace=True) + assert out is d + + def test_curvilinear_grid_passthrough(self): + # A curvilinear (post-'map') grid: every per-axis array already has + # the full nodal shape, not just a 1-D axis. + nx, ny = 3, 2 + gx, gy = np.meshgrid(np.linspace(0.0, 1.0, nx + 1), + np.linspace(0.0, 1.0, ny + 1), + indexing="ij") + d = _make([gx, gy], np.ones((nx, ny, 1))) + out = operations.grid(d) + assert out.get_values().shape == (nx + 1, ny + 1, 2) + np.testing.assert_allclose(out.get_values()[..., 0], gx) + np.testing.assert_allclose(out.get_values()[..., 1], gy) + + def test_dimension_mismatch_raises(self): + d = _make([np.linspace(0.0, 1.0, 5)], np.ones((4, 1))) + d.ctx["cells"] = np.array([4, 4]) # claims 2 dims; grid has 1 axis + with pytest.raises(ValueError, match="dimension"): + operations.grid(d) + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.grid(d) + + +# ========================================================= operations.val2coord +class TestVal2coord: + + def _table(self): + # 5 samples, 3 columns: [x, y0, y1] + return _make([np.arange(5.0)], np.arange(15.0).reshape(5, 3)) + + def test_single_x_multiple_y(self): + group = operations.val2coord(self._table(), x="0", y="1,2") + assert isinstance(group, GDataStateGroup) + assert len(group) == 2 + np.testing.assert_allclose(group[0].get_grid()[0], np.arange(5.0) * 3.0) + np.testing.assert_allclose(group[0].get_values().flatten(), + np.arange(5.0) * 3.0 + 1.0) + + def test_periodic_appends_first_sample(self): + group = operations.val2coord(self._table(), x="0", y="1", periodic=True) + d = group[0] + assert d.get_values().shape[0] == 6 + np.testing.assert_allclose(d.get_values().flatten()[-1], + d.get_values().flatten()[0]) + + def test_mismatched_x_y_counts_raises(self): + with pytest.raises(ValueError): + operations.val2coord(self._table(), x="0,1", y="2") + + def test_colon_range_selector_with_negative_indices_and_step(self): + # 4 columns; "-3:-1:1" exercises the negative-lo, negative-hi, and + # explicit-step branches of the 'lo:hi[:step]' grammar in one shot. + d = _make([np.arange(6.0)], np.arange(24.0).reshape(6, 4)) + group = operations.val2coord(d, x="0", y="-3:-1:1") + assert len(group) == 2 # columns 1, 2 + + @needs_gkeyll + def test_rejects_modal_data(self): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.val2coord(d, x="0", y="1") + + +# ===================================================== operations.extract_input +class TestExtractInput: + + def test_missing_returns_empty_string(self): + d = _make([np.linspace(0.0, 1.0, 3)], np.ones((2, 1))) + assert operations.extract_input(d) == "" + + def test_decodes_base64_ctx_field(self): + text = "title = my sim\nnFrames = 10\n" + encoded = base64.encodebytes(text.encode("utf-8")).decode("utf-8") + d = _make([np.linspace(0.0, 1.0, 3)], np.ones((2, 1)), input_file=encoded) + assert operations.extract_input(d) == text + + def test_returns_a_plain_string_not_a_dataset(self): + d = _make([np.linspace(0.0, 1.0, 3)], np.ones((2, 1))) + assert isinstance(operations.extract_input(d), str) diff --git a/tests/test_operations_fit.py b/tests/test_operations_fit.py new file mode 100644 index 00000000..b57bb740 --- /dev/null +++ b/tests/test_operations_fit.py @@ -0,0 +1,276 @@ +"""Tests for the ``fit`` verb -- model fitting on a dataset's grid.""" + +from __future__ import annotations + +from importlib import import_module +import os +import re + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") + + +def _make(grid, values, **ctx): + d = GDataState(ctx=ctx or None) + d.push(list(grid), values) + return d + + +def _linear_dataset(a=2.0, b=1.0, n=20): + edges = np.linspace(0.0, 1.0, n + 1) + centers = 0.5 * (edges[:-1] + edges[1:]) + y = a * centers + b + return _make([edges], y[:, np.newaxis]), centers + + +def _growth_series(a=1.0, b=0.5, n=60): + edges = np.linspace(0.0, 1.0, n + 1) + centers = 0.5 * (edges[:-1] + edges[1:]) + y = a * np.exp(2.0 * b * centers) + return _make([edges], y[:, np.newaxis]), centers + + +def test_linear_fit_recovers_parameters(): + d, _ = _linear_dataset(a=2.0, b=1.0) + out = operations.fit(d, "linear") + params = out.ctx["fit_params"][0] + np.testing.assert_allclose(params, [2.0, 1.0], atol=1e-8) + assert out.ctx["fit_R2"][0] > 0.999 + + +def test_fitted_curve_matches_evaluated_model(): + d, centers = _linear_dataset(a=3.0, b=-2.0) + out = operations.fit(d, "linear") + expected = 3.0 * centers - 2.0 + np.testing.assert_allclose(out.get_values().flatten(), expected, atol=1e-8) + + +def test_explicit_guess_is_used(): + d, _ = _linear_dataset(a=2.0, b=1.0) + out = operations.fit(d, "linear", guess="1.5,0.5") + np.testing.assert_allclose(out.ctx["fit_params"][0], [2.0, 1.0], atol=1e-6) + + +def test_explicit_guess_as_string_matches_sequence(): + d, _ = _linear_dataset(a=2.0, b=1.0) + out_str = operations.fit(d, "linear", guess="1.0,0.0") + out_seq = operations.fit(d, "linear", guess=[1.0, 0.0]) + np.testing.assert_allclose(out_str.ctx["fit_params"][0], + out_seq.ctx["fit_params"][0]) + + +def test_gaussian_fit_rpn_and_multi_component(): + edges = np.linspace(-5.0, 5.0, 51) + centers = 0.5 * (edges[:-1] + edges[1:]) + y0 = 3.0 * np.exp(-0.5 * (centers / 1.0)**2) + y1 = 5.0 * np.exp(-0.5 * ((centers - 1.0) / 2.0)**2) + d = _make([edges], np.stack([y0, y1], axis=-1)) + out = operations.fit(d, "gaussian") + assert len(out.ctx["fit_params"]) == 2 + np.testing.assert_allclose(out.ctx["fit_params"][0][:2], [3.0, 0.0], + atol=1e-3) + + +def test_wrong_dimensionality_raises(): + d, _ = _linear_dataset() + with pytest.raises(ValueError, match="requires"): + operations.fit(d, "plane") # plane needs 2 spatial dims, data has 1 + + +def test_unknown_fit_type_raises(): + d, _ = _linear_dataset() + with pytest.raises(ValueError): + operations.fit(d, "not_a_real_model_@@") + + +def test_drops_collapsed_axes(): + # A 2nd axis collapsed to a single cell (e.g. after select/integrate). + edges0 = np.linspace(0.0, 1.0, 6) + edges1 = np.linspace(0.0, 1.0, 2) # single cell + centers0 = 0.5 * (edges0[:-1] + edges0[1:]) + y = (2.0 * centers0 + 1.0)[:, np.newaxis, np.newaxis] + d = _make([edges0, edges1], y) + out = operations.fit(d, "linear") + np.testing.assert_allclose(out.ctx["fit_params"][0], [2.0, 1.0], atol=1e-8) + assert out.get_values().ndim == 2 # the collapsed axis was dropped + + +def test_grid_already_cell_centered_needs_no_conversion(): + centers = np.linspace(0.0, 1.0, 20) # matches value count -- not +1 + y = 2.0 * centers + 1.0 + d = _make([centers], y[:, np.newaxis]) + out = operations.fit(d, "linear") + np.testing.assert_allclose(out.ctx["fit_params"][0], [2.0, 1.0], atol=1e-8) + + +def test_plane_fit_2d(): + e0, e1 = np.linspace(0.0, 1.0, 6), np.linspace(0.0, 1.0, 5) + c0, c1 = 0.5 * (e0[:-1] + e0[1:]), 0.5 * (e1[:-1] + e1[1:]) + X, Y = np.meshgrid(c0, c1, indexing="ij") + z = 2.0 * X + 3.0 * Y + 1.0 + d = _make([e0, e1], z[..., np.newaxis]) + out = operations.fit(d, "plane") + np.testing.assert_allclose(out.ctx["fit_params"][0], [2.0, 3.0, 1.0], + atol=1e-6) + + +def test_inplace_and_tag_label(): + d, _ = _linear_dataset() + out = operations.fit(d, "linear", tag="t", label="l", inplace=True) + assert out is d + assert d.get_tag() == "t" + assert d.get_label() == "l" + + +def test_printed_statistics_match_known_linear_residuals(capsys): + x = np.arange(-2.0, 3.0) + # Orthogonal to both the constant and linear terms: the fit remains 2*x+3. + residual = np.array([1.0, -2.0, 2.0, -2.0, 1.0]) + d = _make([x], + np.stack([2 * x + 3 + residual, 2 * x + 3 + 2 * residual], axis=-1)) + operations.fit(d, "linear", print_coeffs=True) + components = capsys.readouterr().out.split(" component ")[1:] + assert len(components) == 2 + for comp, output in enumerate(components): + rss = 14 * (comp + 1)**2 + stats = dict(re.findall(r"^ ([^=]+) = (\S+)", output, re.MULTILINE)) + assert float(stats["R^2"]) == pytest.approx(40 / (40 + rss)) + assert float(stats["RSS"]) == pytest.approx(rss) + assert float(stats["RMSE"]) == pytest.approx(np.sqrt(rss / 5)) + assert float(stats["Residual standard error"]) == pytest.approx( + np.sqrt(rss / 3)) + assert "Samples = 5" in output + assert "Parameters = 2; residual degrees of freedom = 3" in output + assert "x range = [-2, 2]" in output + errors = re.findall(r"\+/- (\S+) \(1-sigma\)", output) + np.testing.assert_allclose([float(error) for error in errors], + np.sqrt([rss / 30, rss / 15]), + rtol=1e-6) + + +def test_printed_window_statistics_exclude_unfitted_tail(capsys): + x = np.arange(8.0) + y = np.array([3., 5., 7., 9., 11., 400., -30., 600.]) + d = _make([x], y[:, None]) + out = operations.fit(d, "linear", window=True, min_n=5, print_coeffs=True) + output = capsys.readouterr().out + assert "Samples = 5 of 8 (leading window)" in output + assert "x range = [0, 4]" in output + stats = dict(re.findall(r"^ ([^=]+) = (\S+)", output, re.MULTILINE)) + assert float(stats["R^2"]) == pytest.approx(1.0) + assert float(stats["RMSE"]) == pytest.approx(0.0, abs=1e-8) + assert out.values.shape == d.values.shape + + +def test_printed_statistics_explain_undefined_estimates(capsys): + d = _make([np.array([0., 1.])], np.array([[3.], [3.]])) + operations.fit(d, "linear", guess=[0., 3.], print_coeffs=True) + output = capsys.readouterr().out + assert "R^2 = undefined (constant fitted data)" in output + assert "Residual standard error = undefined" in output + assert "no residual degrees of freedom" in output + assert "1-sigma uncertainty unavailable" in output + + +@needs_gkeyll +def test_rejects_modal_data(): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.fit(d, "linear") + + +# ── window=True -- growth-rate-style leading-window fits ───────────────────── + + +def test_window_recovers_growth_rate(): + d, _ = _growth_series(a=1.0, b=1.5) + out = operations.fit(d, "exp2", window=True) + assert out.ctx["fit_params"][0][1] == pytest.approx(1.5, abs=1e-3) + + +def test_window_output_shape_matches_full_grid(): + d, centers = _growth_series() + out = operations.fit(d, "exp2", window=True) + assert out.get_values().shape[0] == len(centers) + + +def test_window_explicit_guess_string_and_sequence_agree(): + d, _ = _growth_series(a=1.0, b=0.8) + out_str = operations.fit(d, "exp2", window=True, guess="1,1") + out_seq = operations.fit(d, "exp2", window=True, guess=(1.0, 1.0)) + np.testing.assert_allclose(out_str.ctx["fit_params"][0], + out_seq.ctx["fit_params"][0]) + + +def test_window_min_n_controls_minimum_window(): + d, _ = _growth_series(a=1.0, b=1.0, n=100) + out = operations.fit(d, "exp2", window=True, min_n=5) + assert out.ctx["fit_params"][0][1] == pytest.approx(1.0, abs=1e-2) + + +def test_window_inplace_and_tag_label(): + d, _ = _growth_series() + out = operations.fit(d, + "exp2", + window=True, + tag="g", + label="growth-fit", + inplace=True) + assert out is d + assert d.get_tag() == "g" + assert d.get_label() == "growth-fit" + + +def test_window_rejects_multi_dim_data(): + e0, e1 = np.linspace(0.0, 1.0, 6), np.linspace(0.0, 1.0, 5) + c0, c1 = 0.5 * (e0[:-1] + e0[1:]), 0.5 * (e1[:-1] + e1[1:]) + X, Y = np.meshgrid(c0, c1, indexing="ij") + d = _make([e0, e1], (X + Y)[..., np.newaxis]) + with pytest.raises(ValueError, match="window=True is only supported"): + operations.fit(d, "plane", window=True) + + +@needs_gkeyll +def test_window_rejects_modal_data(): + d = pg.load(F1) + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.fit(d, "exp2", window=True) + + +def test_growth_is_the_declared_leading_window_composition(monkeypatch): + growth_module = import_module("postgkyl.operations.growth") + calls = [] + + def fake_fit(*args, **kwargs): + calls.append((args, kwargs)) + return "result" + + monkeypatch.setattr(growth_module, "fit", fake_fit) + data = object() + result = growth_module.growth(data, + guess="1,2", + min_n=7, + inplace=True, + tag="fit", + label="growth") + assert result == "result" + assert calls == [((data, "exp2"), { + "guess": "1,2", + "window": True, + "min_n": 7, + "inplace": True, + "tag": "fit", + "label": "growth", + })] diff --git a/tests/test_operations_gk_fluxsurf.py b/tests/test_operations_gk_fluxsurf.py new file mode 100644 index 00000000..243572d9 --- /dev/null +++ b/tests/test_operations_gk_fluxsurf.py @@ -0,0 +1,173 @@ +"""Characterization tests for the moved gyrokinetic flux-surface operation.""" + +from __future__ import annotations + +from dataclasses import replace +from importlib import import_module +import os +from types import SimpleNamespace + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.operations import gyrokinetics as gk_ops + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +FIELD = os.path.join(ROOT, "tests", "test_data", + "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +@needs_gkeyll +def test_flux_surface_move_preserves_output_and_projection_reuse(): + data = pg.load(FIELD) + geometry = gk_ops.resolve_geometry(data.file_name) + grid = gk_ops.resolve_flux_surface_grid(data, + geometry, + x_idx=0, + nphi=4, + nz_interp=2) + first = gk_ops.extract_flux_surface(data, grid) + second = gk_ops.extract_flux_surface(data.clone(), grid) + assert first.values.shape == (4, 64, 1) + assert first.ctx["interpolated"] is True + np.testing.assert_allclose(first.values, second.values) + + +@needs_gkeyll +@pytest.mark.parametrize(("kwargs", "message"), [ + ({ + "nphi": 0 + }, "nphi must be a positive integer"), + ({ + "nz_interp": 0 + }, "nz_interp must be a positive integer"), + ({ + "x_idx": -1 + }, "out of bounds"), +]) +def test_flux_surface_public_validation(kwargs, message): + data = pg.load(FIELD) + geometry = gk_ops.resolve_geometry(data.file_name) + with pytest.raises(ValueError, match=message): + gk_ops.resolve_flux_surface_grid(data, geometry, **kwargs) + + +@needs_gkeyll +def test_flux_surface_grid_requires_toroidal_geometry_and_integer_index(): + data = pg.load(FIELD) + geometry = gk_ops.resolve_geometry(data.file_name) + with pytest.raises(ValueError, match="no toroidal-angle component"): + gk_ops.resolve_flux_surface_grid(data, replace(geometry, phi=None)) + with pytest.raises(ValueError, match="x_idx must be an integer"): + gk_ops.resolve_flux_surface_grid(data, geometry, x_idx=True) + + +@needs_gkeyll +def test_flux_surface_grid_requires_two_binormal_and_parallel_points(): + data = pg.load(FIELD).clone() + data.ctx["poly_order"] = 0 + data._grid[1] = np.array([0.0, 1.0]) + geometry = gk_ops.resolve_geometry(data.file_name) + with pytest.raises(ValueError, match="at least two interpolated y and z"): + gk_ops.resolve_flux_surface_grid(data, geometry) + + +@needs_gkeyll +def test_extract_flux_surface_validates_reusable_grid_metadata(): + data = pg.load(FIELD) + geometry = gk_ops.resolve_geometry(data.file_name) + grid = gk_ops.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) + + shifted = data.clone() + shifted.grid[0] = shifted.grid[0] + 0.1 + with pytest.raises(ValueError, match="computational grid does not match"): + gk_ops.extract_flux_surface(shifted, grid) + + with pytest.raises(ValueError, match="out of bounds"): + gk_ops.extract_flux_surface(data, replace(grid, x_idx=10_000)) + + with pytest.raises(ValueError, match="projection and data grid shapes"): + gk_ops.extract_flux_surface(data, replace(grid, phi_2d=np.ones((1, 1)))) + + +@needs_gkeyll +def test_extract_flux_surface_rejects_zero_toroidal_span(): + data = pg.load(FIELD) + geometry = gk_ops.resolve_geometry(data.file_name) + grid = gk_ops.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) + zero_span = replace(grid, phi_2d=np.zeros_like(grid.phi_2d)) + with pytest.raises(ValueError, match="zero or non-finite"): + gk_ops.extract_flux_surface(data, zero_span) + + +def test_flux_surface_grid_collection_caches_by_geometry_prefix(monkeypatch): + fluxsurf = import_module("postgkyl.operations.gyrokinetics.fluxsurf") + first = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + repeated = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + second = SimpleNamespace(file_name="block-two", ctx={"block": 2}) + calls = [] + monkeypatch.setattr(fluxsurf, "geometry_prefix", lambda path: path) + monkeypatch.setattr( + fluxsurf, "resolve_geometry", lambda path, **kwargs: calls.append( + (path, kwargs)) or path) + monkeypatch.setattr(fluxsurf, "resolve_flux_surface_grid", + lambda data, geo, **_kwargs: f"grid:{geo}") + + grids = fluxsurf.flux_surface_grids([first, repeated, second], + mapc2p="map-*.gkyl", + nodes_file="nodes-*.gkyl") + assert grids == { + "block-one": "grid:block-one", + "block-two": "grid:block-two", + } + assert calls == [ + ("block-one", { + "mapc2p": "map-1.gkyl", + "nodes_file": "nodes-1.gkyl" + }), + ("block-two", { + "mapc2p": "map-2.gkyl", + "nodes_file": "nodes-2.gkyl" + }), + ] + assert fluxsurf.grid_for(grids, first) == "grid:block-one" + + +def test_gk_fluxsurf_composes_geometry_grid_and_extraction(monkeypatch): + fluxsurf = import_module("postgkyl.operations.gyrokinetics.fluxsurf") + data = SimpleNamespace(file_name="field.gkyl") + calls = [] + monkeypatch.setattr( + fluxsurf, "resolve_geometry", lambda path, **kwargs: calls.append( + ("geometry", path, kwargs)) or "geo") + monkeypatch.setattr( + fluxsurf, "resolve_flux_surface_grid", + lambda source, geo, **kwargs: calls.append( + ("grid", source, geo, kwargs)) or "grid") + monkeypatch.setattr( + fluxsurf, "extract_flux_surface", + lambda source, grid, **kwargs: calls.append( + ("extract", source, grid, kwargs)) or "result") + + result = fluxsurf.gk_fluxsurf(data, + mapc2p="map.gkyl", + x_idx=2, + nphi=16, + nz_interp=3, + comp=4, + inplace=True, + tag="surface", + label="flux") + assert result == "result" + assert [call[0] for call in calls] == ["geometry", "grid", "extract"] + assert calls[-1][-1] == { + "comp": 4, + "inplace": True, + "tag": "surface", + "label": "flux" + } diff --git a/tests/test_operations_gk_geometry.py b/tests/test_operations_gk_geometry.py new file mode 100644 index 00000000..e12324cd --- /dev/null +++ b/tests/test_operations_gk_geometry.py @@ -0,0 +1,210 @@ +"""Unit contracts for shared gyrokinetic geometry machinery.""" + +from __future__ import annotations + +from dataclasses import replace +from importlib import import_module +from types import SimpleNamespace + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython + +geometry = import_module("postgkyl.operations.gyrokinetics.geometry") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +def _valid_geometry(num_dims=2, *, phi=False, corner=None): + coords = [np.array([0.0, 1.0]) for _ in range(num_dims)] + shape = (2, ) * num_dims + values = np.ones(shape) + return geometry.Geometry(coords, values, 2.0 * values, + values if phi else None, corner) + + +def test_gauss_nodes_are_ordered_inside_each_cell(): + nodes = geometry._gauss_nodes(np.array([0.0, 2.0, 4.0])) + assert nodes.shape == (4, ) + assert np.all(np.diff(nodes) > 0.0) + assert nodes[0] > 0.0 and nodes[-1] < 4.0 + + +def test_pointwise_file_squeezes_grid_and_values(monkeypatch): + state = SimpleNamespace(grid=[np.array([[0.0, 1.0, 2.0]])], + values=np.ones((1, 2, 1))) + monkeypatch.setattr(geometry, "GDataState", lambda _path: state) + grid, values, returned = geometry._pointwise_file("geometry.gkyl") + np.testing.assert_array_equal(grid[0], [0.0, 1.0, 2.0]) + np.testing.assert_array_equal(values, [1.0, 1.0]) + assert returned is state + + +def test_geometry_components_support_cartesian_and_rz_layouts(): + mapc2p = SimpleNamespace(ctx={"geometry_type": geometry._MAPC2P_IDX}) + values = np.array([[[3.0, 4.0, 2.0], [0.0, 2.0, 5.0]]]) + major_r, vert_z, phi = geometry._geometry_components(values, mapc2p, + "map.gkyl") + np.testing.assert_allclose(major_r, [[5.0, 2.0]]) + np.testing.assert_allclose(vert_z, [[2.0, 5.0]]) + assert phi.shape == (1, 2) + + rz = SimpleNamespace(ctx={"geometry_type": 1}) + values_3d = np.ones((2, 2, 2, 3)) + major_r, vert_z, phi = geometry._geometry_components(values_3d, rz, "rz.gkyl") + assert major_r.shape == vert_z.shape == phi.shape == (2, 2, 2) + + +@pytest.mark.parametrize(("ctx", "values", "message"), [ + ({ + "geometry_type": geometry._MAPC2P_IDX + }, np.ones((2, 2, 2)), "at least 3 Cartesian"), + ({ + "geometry_type": 1 + }, np.ones(2), "at least 2 R/Z"), +]) +def test_geometry_components_reject_short_layouts(ctx, values, message): + with pytest.raises(ValueError, match=message): + geometry._geometry_components(values, SimpleNamespace(ctx=ctx), "bad.gkyl") + + +def test_read_nodes_geometry_recovers_gauss_coordinates(monkeypatch): + grid = [np.array([0.0, 0.5, 1.0]), np.array([-1.0, 0.0, 1.0])] + values = np.ones((2, 2, 3)) + data = SimpleNamespace(ctx={"geometry_type": geometry._MAPC2P_IDX}) + monkeypatch.setattr(geometry, "_pointwise_file", lambda _path: + (grid, values, data)) + coords, major_r, vert_z, phi = geometry._read_nodes_geometry("nodes.gkyl") + assert [axis.shape for axis in coords] == [(2, ), (2, )] + assert major_r.shape == vert_z.shape == phi.shape == (2, 2) + + +def test_read_nodes_geometry_rejects_unknown_layout(monkeypatch): + grid = [np.array([0.0, 1.0])] + values = np.ones((2, 3)) + data = SimpleNamespace(ctx={"geometry_type": geometry._MAPC2P_IDX}) + monkeypatch.setattr(geometry, "_pointwise_file", lambda _path: + (grid, values, data)) + with pytest.raises(ValueError, match="Unrecognized nodal geometry layout"): + geometry._read_nodes_geometry("nodes.gkyl") + + +def test_read_corner_geometry_builds_point_coordinates(monkeypatch): + grid = [np.array([0.0, 3.0]), np.array([-2.0, 2.0])] + values = np.ones((3, 4, 3)) + data = SimpleNamespace(ctx={"geometry_type": geometry._MAPC2P_IDX}) + monkeypatch.setattr(geometry, "_pointwise_file", lambda _path: + (grid, values, data)) + coords, major_r, vert_z = geometry._read_corner_rz("corner.gkyl") + assert [axis.size for axis in coords] == [3, 4] + assert major_r.shape == vert_z.shape == (3, 4) + + +@pytest.mark.parametrize(("candidate", "num_dims", "message"), [ + (_valid_geometry(1), 2, "Geometry has 1 dimensions"), + (geometry.Geometry([np.array([[0.0, 1.0]])], np.ones((2, )), np.ones( + (2, )), None, None), 1, "one-dimensional arrays"), + (geometry.Geometry([np.array([0.0, 1.0, 0.5])], np.ones( + (3, )), np.ones((3, )), None, None), 1, "strictly monotonic"), + (geometry.Geometry([np.array([0.0, 1.0])], np.ones((3, )), np.ones( + (2, )), None, None), 1, "R/Z array shapes are incompatible"), + (replace(_valid_geometry(2), phi=np.ones( + (2, ))), 2, "toroidal-angle shape"), + (_valid_geometry(2, + corner=([np.array([0.0, 1.0])], np.ones( + (2, )), np.ones( + (2, )))), 2, "Corner geometry has 1 dimensions"), + (_valid_geometry( + 2, + corner=([np.array([0.0]), np.array([0.0, 1.0])], np.ones( + (1, 2)), np.ones( + (1, 2)))), 2, "Corner geometry coordinate and R/Z"), +]) +def test_validate_geometry_rejects_each_shape_invariant(candidate, num_dims, + message): + with pytest.raises(ValueError, match=message): + geometry._validate_geometry(candidate, num_dims) + + +def test_resolve_geometry_honors_explicit_nodes_and_loads_corner( + monkeypatch, tmp_path): + source = tmp_path / "sim-field_0.gkyl" + nodes = tmp_path / "nodes.gkyl" + corner = tmp_path / "sim-geo_corn_nodes.gkyl" + coords = [np.array([0.0, 1.0]), np.array([-1.0, 1.0])] + values = np.ones((2, 2)) + calls = [] + monkeypatch.setattr(geometry.os.path, "exists", + lambda path: path in {str(nodes), str(corner)}) + monkeypatch.setattr( + geometry, "_read_nodes_geometry", lambda path: + (calls.append(path) or (coords, values, values, None))) + monkeypatch.setattr( + geometry, "_read_corner_rz", lambda path: + (calls.append(path) or (coords, values, values))) + resolved = geometry.resolve_geometry(str(source), nodes_file=str(nodes)) + assert calls == [str(nodes), str(corner)] + assert resolved.corner is not None + + +def test_resolve_geometry_without_a_name_requires_an_override(): + with pytest.raises(ValueError, match="Could not find a geometry file"): + geometry.resolve_geometry(None) + + +def test_validate_modal_data_reports_missing_data_and_metadata(): + empty = pg.GData() + with pytest.raises(ValueError, match="loaded dataset"): + geometry._validate_modal_data(empty, "projection", (0, )) + + no_basis = pg.GData() + no_basis.push([np.array([0.0, 1.0])], np.ones((1, 1))) + with pytest.raises(ValueError, match="basis_type"): + geometry._validate_modal_data(no_basis, "projection", (1, )) + + no_basis.ctx["basis_type"] = "serendipity" + no_basis.ctx["poly_order"] = True + with pytest.raises(ValueError, match="nonnegative integer"): + geometry._validate_modal_data(no_basis, "projection", (1, )) + + +def test_validate_modal_data_reports_grid_shape_and_monotonicity(): + data = pg.GData(ctx={ + "basis_type": "serendipity", + "poly_order": 0, + "value_form": "modal", + }) + data.push([np.array([0.0, 1.0])], np.ones((1, 1))) + data._grid = [np.array([0.0])] + with pytest.raises(ValueError, match="one-dimensional edge grid"): + geometry._validate_modal_data(data, "projection", (1, )) + + data._grid = [np.array([0.0, 1.0, 0.5])] + with pytest.raises(ValueError, match="strictly monotonic"): + geometry._validate_modal_data(data, "projection", (1, )) + + +@needs_gkeyll +def test_num_fields_rejects_incompatible_coefficient_count(): + data = pg.GData(ctx={ + "basis_type": "serendipity", + "poly_order": 1, + "value_form": "modal", + }) + data.push([np.array([0.0, 1.0])], np.ones((1, 3))) + with pytest.raises(ValueError, match="incompatible"): + geometry._num_fields(data) + + +def test_validate_component_rejects_boolean_before_basis_lookup(): + with pytest.raises(ValueError, match="integer component"): + geometry._validate_component(pg.GData(), True) + + +def test_same_grid_rejects_dimension_and_shape_mismatches(): + axis = np.array([0.0, 1.0]) + assert not geometry._same_grid([axis], [axis, axis]) + assert not geometry._same_grid([axis], [np.array([0.0, 0.5, 1.0])]) diff --git a/tests/test_operations_gk_rz.py b/tests/test_operations_gk_rz.py new file mode 100644 index 00000000..ba50553f --- /dev/null +++ b/tests/test_operations_gk_rz.py @@ -0,0 +1,308 @@ +"""Gyrokinetic R-Z operation, public surfaces, and compatibility paths.""" + +from __future__ import annotations + +from dataclasses import replace +from importlib import import_module +import os +from types import SimpleNamespace + +import click +import numpy as np +import pytest +from click.testing import CliRunner + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.cli.app import COMMANDS + +gk_rz_command = next(command for command in COMMANDS if command.name == "gk_rz") +from postgkyl.cli.state import DataSpace +from postgkyl.operations import gyrokinetics as gk_ops + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1D = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F2D = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-elc_M2par_10.gkyl") +F2D_GEO = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-geo_int_mapc2p.gkyl") +F3D = os.path.join(DATA, "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +@needs_gkeyll +def test_2d_mapping_reference_grid_and_values(): + mapped = pg.gk_rz(pg.load(F2D), nz_interp=2) + assert [axis.shape for axis in mapped.grid] == [(33, 33), (33, 33)] + assert mapped.values.shape == (32, 32, 1) + np.testing.assert_allclose(mapped.values.flat[:5], [ + 3.12774618e30, 3.15719364e30, 3.23307916e30, 3.18034806e30, 3.15422838e30 + ], + rtol=2e-9) + + +@needs_gkeyll +def test_3d_mapping_reference_and_fft_phase(): + data = pg.load(F3D) + geometry = gk_ops.resolve_geometry(data.file_name) + projection = gk_ops.resolve_rz_projection(data, geometry, nz_interp=2) + at_zero = gk_ops.map_to_rz(data, projection, phi_tor=0.0) + at_quarter = gk_ops.map_to_rz(data, projection, phi_tor=np.pi / 2) + assert [axis.shape for axis in at_zero.grid] == [(97, 65), (97, 65)] + assert at_zero.values.shape == (96, 64, 1) + np.testing.assert_allclose(at_zero.values.flat[:5], [ + 9.97245134e18, 9.71438453e18, 9.57553863e18, 9.50610482e18, 9.81316530e18 + ], + rtol=2e-9) + assert not np.allclose(at_zero.values, at_quarter.values) + + +def test_geometry_prefers_nodes_and_honors_explicit_modal_override( + tmp_path, monkeypatch): + from postgkyl.operations.gyrokinetics import geometry as geometry_module + + source = tmp_path / "sim-field_0.gkyl" + nodes = tmp_path / "sim-geo_int_nodes.gkyl" + modal = tmp_path / "sim-geo_int_mapc2p.gkyl" + nodes.touch() + modal.touch() + coords = [np.array([0.0, 1.0]), np.array([-1.0, 1.0])] + arrays = np.ones((2, 2)) + calls = [] + monkeypatch.setattr( + geometry_module, "_read_nodes_geometry", lambda path: (calls.append( + ("nodes", path)) or (coords, arrays, arrays, None))) + monkeypatch.setattr( + geometry_module, "_read_mapc2p_geometry", lambda path: (calls.append( + ("mapc2p", path)) or (coords, arrays, arrays, None))) + + gk_ops.resolve_geometry(str(source)) + assert calls[-1] == ("nodes", str(nodes)) + gk_ops.resolve_geometry(str(source), mapc2p="") + assert calls[-1] == ("mapc2p", str(modal)) + + +@needs_gkeyll +def test_geometry_overrides_and_validation_errors(tmp_path): + data = pg.load(F2D) + with pytest.raises(ValueError, match="either mapc2p=.*nodes_file"): + pg.gk_rz(data, mapc2p=F2D_GEO, nodes_file=F2D_GEO) + explicit = pg.gk_rz(data, mapc2p=F2D_GEO, nz_interp=2) + inferred = pg.gk_rz(data, nz_interp=2) + np.testing.assert_allclose(explicit.values, inferred.values) + + missing = data.clone() + missing._file_name = str(tmp_path / "absent-field_0.gkyl") + with pytest.raises(ValueError, match="Could not find a geometry file"): + pg.gk_rz(missing) + with pytest.raises(ValueError, match="positive integer"): + pg.gk_rz(data, nz_interp=0) + with pytest.raises(ValueError, match="out of bounds"): + pg.gk_rz(data, comp=1) + with pytest.raises(ValueError, match="requires 2-D or 3-D"): + pg.gk_rz(pg.load(F1D)) + with pytest.raises(ValueError, match="un-interpolated modal DG"): + pg.gk_rz(data.interpolate()) + + +@needs_gkeyll +def test_missing_toroidal_geometry_and_incompatible_projection_fail_clearly(): + data = pg.load(F3D) + coords = [np.array([0.0, 1.0])] * 3 + values = np.ones((2, 2, 2)) + no_phi = gk_ops.Geometry(coords=coords, + major_r=values, + vert_z=values, + phi=None, + corner=None) + with pytest.raises(ValueError, match="no toroidal-angle component"): + gk_ops.resolve_rz_projection(data, no_phi) + + geometry = gk_ops.resolve_geometry(data.file_name) + projection = gk_ops.resolve_rz_projection(data, geometry, nz_interp=2) + shifted = data.clone() + shifted.grid[0] = shifted.grid[0] + 0.01 + with pytest.raises(ValueError, match="computational grid does not match"): + gk_ops.map_to_rz(shifted, projection) + + +@needs_gkeyll +def test_3d_projection_rejects_thin_and_zero_span_geometry(): + data = pg.load(F3D) + geometry = gk_ops.resolve_geometry(data.file_name) + + thin = data.clone() + thin.ctx["poly_order"] = 0 + thin._grid[1] = np.array([0.0, 1.0]) + with pytest.raises(ValueError, match="at least two interpolated y and z"): + gk_ops.resolve_rz_projection(thin, geometry) + + zero_span = replace(geometry, phi=np.zeros_like(geometry.phi)) + with pytest.raises(ValueError, match="zero or non-finite"): + gk_ops.resolve_rz_projection(data, zero_span) + + +@needs_gkeyll +def test_3d_projection_uses_corner_geometry_when_available(): + data = pg.load(F3D) + geometry = gk_ops.resolve_geometry(data.file_name) + corner_coords = [np.array(axis, copy=True) for axis in geometry.coords] + dz = geometry.coords[2][-1] - geometry.coords[2][0] + corner_coords[2] = np.array( + [geometry.coords[2][0] - dz, geometry.coords[2][-1] + dz]) + corner_r = np.stack([geometry.major_r[..., 0], geometry.major_r[..., -1]], + axis=-1) + corner_z = np.stack([geometry.vert_z[..., 0], geometry.vert_z[..., -1]], + axis=-1) + with_corner = replace(geometry, corner=(corner_coords, corner_r, corner_z)) + projection = gk_ops.resolve_rz_projection(data, with_corner, nz_interp=2) + assert projection.r.shape == projection.z.shape == (97, 65) + + +@needs_gkeyll +def test_reusable_rz_projection_validates_every_shape_contract(): + data_2d = pg.load(F2D) + geometry_2d = gk_ops.resolve_geometry(data_2d.file_name) + projection_2d = gk_ops.resolve_rz_projection(data_2d, geometry_2d) + + invalid_2d = [ + (replace(projection_2d, num_dims=4), "invalid dimensionality"), + (replace(projection_2d, num_dims=3), "dimensionality does not match"), + (replace(projection_2d, + z=projection_2d.z[:, :-1]), "matching 2-D arrays"), + (replace(projection_2d, r=projection_2d.r[:-1], + z=projection_2d.z[:-1]), "expected grid shape"), + ] + for projection, message in invalid_2d: + with pytest.raises(ValueError, match=message): + gk_ops.map_to_rz(data_2d, projection) + + data_3d = pg.load(F3D) + geometry_3d = gk_ops.resolve_geometry(data_3d.file_name) + projection_3d = gk_ops.resolve_rz_projection(data_3d, + geometry_3d, + nz_interp=2) + invalid_3d = [ + (replace(projection_3d, zc=None), "metadata is incomplete"), + (replace(projection_3d, box=0.0), "span must be finite and nonzero"), + (replace(projection_3d, wind=projection_3d.wind[:-1]), + "projection and data grid shapes differ"), + ] + for projection, message in invalid_3d: + with pytest.raises(ValueError, match=message): + gk_ops.map_to_rz(data_3d, projection) + + +def test_rz_projection_collection_caches_by_geometry_prefix(monkeypatch): + rz = import_module("postgkyl.operations.gyrokinetics.rz") + first = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + repeated = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + second = SimpleNamespace(file_name="block-two", ctx={"block": 2}) + calls = [] + monkeypatch.setattr(rz, "geometry_prefix", lambda path: path) + monkeypatch.setattr( + rz, "resolve_geometry", lambda path, **kwargs: calls.append( + (path, kwargs)) or path) + monkeypatch.setattr(rz, "resolve_rz_projection", + lambda data, geo, **_kwargs: f"projection:{geo}") + + projections = rz.rz_projections([first, repeated, second], + mapc2p="map-*.gkyl", + nodes_file="nodes-*.gkyl") + assert projections == { + "block-one": "projection:block-one", + "block-two": "projection:block-two", + } + assert calls == [ + ("block-one", { + "mapc2p": "map-1.gkyl", + "nodes_file": "nodes-1.gkyl" + }), + ("block-two", { + "mapc2p": "map-2.gkyl", + "nodes_file": "nodes-2.gkyl" + }), + ] + assert rz.projection_for(projections, first) == "projection:block-one" + + +@needs_gkeyll +def test_state_propagation_projection_reuse_and_public_surfaces(): + + class DerivedData(pg.GData): + pass + + source = DerivedData(F2D, tag="source", label="original") + original = source.values.copy() + geometry = gk_ops.resolve_geometry(source.file_name) + projection = gk_ops.resolve_rz_projection(source, geometry, nz_interp=2) + first = gk_ops.map_to_rz(source, projection, tag="rz", label="mapped") + second = gk_ops.map_to_rz(source.clone(), projection) + assert isinstance(first, DerivedData) + assert first is not source + assert first.file_name == source.file_name + assert (first.tag, first.label, first.ctx["interpolated"]) == ("rz", "mapped", + True) + np.testing.assert_array_equal(source.values, original) + np.testing.assert_allclose(first.values, second.values) + + fluent = source.gk_rz(mapc2p=F2D_GEO, nz_interp=2) + functional = pg.gk_rz(source, mapc2p=F2D_GEO, nz_interp=2) + np.testing.assert_allclose(fluent.values, functional.values) + assert pg.gk_rz is gk_ops.gk_rz + + inplace = source.clone() + result = pg.gk_rz(inplace, mapc2p=F2D_GEO, nz_interp=2, inplace=True) + assert result is inplace and result.ctx["interpolated"] is True + + +@needs_gkeyll +def test_comp_selects_an_explicit_physical_field(): + source = pg.load(F2D) + multi = pg.GData(ctx={ + key: value + for key, value in source.ctx.items() if key != "num_comps" + }) + multi.push([axis.copy() for axis in source.grid], + np.concatenate([source.values, 2.0 * source.values], axis=-1)) + multi._file_name = source.file_name + first = pg.gk_rz(multi, comp=0, nz_interp=2) + second = pg.gk_rz(multi, comp=1, nz_interp=2) + np.testing.assert_allclose(second.values, 2.0 * first.values) + + +@needs_gkeyll +def test_group_compatibility_cli_and_help_section(): + from postgkyl.cli.app import cli + from postgkyl.diagnostics.gk import fluxsurf as old_fluxsurf + from postgkyl.diagnostics.gk import rz as old_rz + + group = pg.GDataGroup([pg.load(F2D), pg.load(F2D)]) + mapped = group.gk_rz(mapc2p=F2D_GEO, nz_interp=2) + assert isinstance(mapped, pg.GDataGroup) and len(mapped) == 2 + + assert old_rz.gk_rz is gk_ops.gk_rz + assert old_rz.RzProjection is gk_ops.RzProjection + assert old_fluxsurf.FluxSurfaceGrid is gk_ops.FluxSurfaceGrid + assert old_fluxsurf.extract_flux_surface is gk_ops.extract_flux_surface + + space = DataSpace(datasets=[pg.load(F2D)]) + with click.Context(gk_rz_command, obj=space) as ctx: + ctx.invoke(gk_rz_command, + mapc2p=F2D_GEO, + nodes_file=None, + z_axis=0.0, + phi_tor=0.0, + nz_interp=2, + use=None, + tag="rz", + label=None) + expected = pg.gk_rz(pg.load(F2D), mapc2p=F2D_GEO, nz_interp=2, tag="rz") + np.testing.assert_allclose(space.datasets[0].values, expected.values) + + help_text = CliRunner().invoke(cli, ["--help"]).output + verbs = help_text.split("Diagnostics:", 1)[0] + diagnostics = help_text.split("Diagnostics:", 1)[1].split("Render:", 1)[0] + assert "gk_rz" in verbs and "gk_rz" not in diagnostics diff --git a/tests/test_operations_local_poly.py b/tests/test_operations_local_poly.py new file mode 100644 index 00000000..cd85af0f --- /dev/null +++ b/tests/test_operations_local_poly.py @@ -0,0 +1,134 @@ +"""Tests for the ``local_poly`` verb -- the discontinuity-preserving DG +plotting mesh (see ``dg.interpolate.local_poly``), ported from the old +Typer-era ``dg_local_poly`` command (``PLAN.md``/``14-cli.md``: initially +deferred, since the old implementation depended on hand-derived per-order +polynomial tables (``modalDG/kernels/expand_*d.py``, serendipity only) with +no equivalent in the new engine -- superseded here by +``gpython.basis.eval_matrix``, which evaluates any basis at arbitrary +points through Gkeyll's own compiled basis-eval). +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") +pytestmark = needs_gkeyll + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GEN = os.path.join(DATA, "generated") +F1D = os.path.join(GEN, "1d_ms_p1.gkyl") +F2D = os.path.join(GEN, "2d_ms_p2.gkyl") +F3D = os.path.join(GEN, "3d_ms_p1.gkyl") + +# 1D p1 serendipity basis, evaluated at the reference-cell faces. +_B0 = 0.7071067811865475 +_B1 = 1.224744871391589 + + +def test_matches_hand_evaluated_basis_at_cell_faces(): + d = pg.load(F1D) + c0, c1 = d.get_values()[0] + expect_left = _B0 * c0 - _B1 * c1 + expect_right = _B0 * c0 + _B1 * c1 + + lp = d.local_poly(npoints=2) + np.testing.assert_allclose(lp.get_values()[0, 0], expect_left) + np.testing.assert_allclose(lp.get_values()[1, 0], expect_right) + assert lp.grid[0][0] == pytest.approx(0.0) + assert lp.grid[0][1] == pytest.approx(1.0 / d.num_cells[0]) + + +def test_inserts_nan_at_every_interior_cell_interface(): + d = pg.load(F1D) + npoints = 3 + lp = d.local_poly(npoints=npoints) + values = lp.get_values()[:, 0] + num_cells = int(d.num_cells[0]) + + # One NaN spliced in after every cell's block of `npoints` raw points, + # except the last cell (no interface after the final face). + assert values.shape[0] == npoints * num_cells + (num_cells - 1) + nan_positions = np.flatnonzero(np.isnan(values)) + # np.insert's k-th insertion index shifts by k in the resulting array + # (each prior insertion pushes it one further along). + pre_insertion = np.arange(npoints, npoints * num_cells, npoints) + expected_positions = pre_insertion + np.arange(len(pre_insertion)) + np.testing.assert_array_equal(nan_positions, expected_positions) + + # The grid coordinate at a NaN repeats the preceding (cell-right-face) + # coordinate, so plotting breaks the line without leaving a coordinate gap. + for pos in nan_positions: + assert lp.grid[0][pos] == pytest.approx(lp.grid[0][pos - 1]) + + +def test_backend_and_flags_after_local_poly(): + lp = pg.load(F1D).local_poly() + assert lp.backend == "numpy" + assert lp.is_interpolated + assert lp.ctx["interpolated"] is True + + +def test_default_npoints_is_two(): + d = pg.load(F1D) + num_cells = int(d.num_cells[0]) + lp = d.local_poly() + # 2 raw points/cell + one NaN at each of the (num_cells - 1) interior faces. + assert lp.get_values().shape[0] == 2 * num_cells + (num_cells - 1) + + +def test_2d_and_3d_shapes(): + d2 = pg.load(F2D) + lp2 = d2.local_poly(npoints=4) + nx, ny = (int(c) for c in d2.num_cells) + assert lp2.grid[0].shape == (4 * nx + (nx - 1), ) + assert lp2.grid[1].shape == (4 * ny + (ny - 1), ) + # `d2.num_comps` counts raw modal coefficients (fields * num_basis); this + # fixture holds a single field. + assert lp2.get_values().shape == (4 * nx + (nx - 1), 4 * ny + (ny - 1), 1) + + d3 = pg.load(F3D) + lp3 = d3.local_poly() + assert lp3.num_dims == 3 + assert not np.all(np.isnan(lp3.get_values())) + + +def test_missing_poly_order_raises(): + d = pg.load(F1D) + del d.ctx["poly_order"] + with pytest.raises(ValueError, match="poly_order"): + d.local_poly() + + +def test_missing_basis_type_raises(): + d = pg.load(F1D) + del d.ctx["basis_type"] + with pytest.raises(ValueError, match="basis_type"): + d.local_poly() + + +def test_rejects_non_modal_value_form(): + d = pg.load(F1D).to_nodal() + with pytest.raises(ValueError, match="modal value_form"): + d.local_poly() + + +def test_inplace_and_tag_label(): + d = pg.load(F1D) + same = d.local_poly(inplace=True, tag="lp", label="lp-label") + assert same is d + assert d.tag == "lp" + assert d.get_label() == "lp-label" + + d2 = pg.load(F1D) + new = d2.local_poly(tag="lp2") + assert new is not d2 + assert new.tag == "lp2" diff --git a/tests/test_operations_map.py b/tests/test_operations_map.py new file mode 100644 index 00000000..a591b0d4 --- /dev/null +++ b/tests/test_operations_map.py @@ -0,0 +1,418 @@ +"""Tests for the ``map`` verb (grid mapping) and the ``select`` curvilinear +guard it motivates. See ``MAPPING.md`` for the design; ``postgkyl.dg.map`` is +the (already-tested, layer-03) engine this verb delegates to. + +Mapping fields are built two ways: + +- **synthetically** (``_synthetic_map``/``_project_1d``/``_project_2d``, + mirroring ``tests/test_dg_map.py``): exact per-cell coefficients of a + chosen physical-coordinate function, so the expected grid is computable + independently of the code under test. +- **from the real generated fixtures** (``generated/2d_c2p_*.gkyl``) for a + genuine file-based conf-space integration test. + +**Conf-space maps (``mapc2p``/``mc2nu``) are one joint ``m``-D curvilinear +map**: every physical coordinate is evaluated over all ``m`` mapped +dimensions, so a non-separable map (e.g. a rotation) is representable. +**Velocity-space maps (``mapc2p_vel``) are diagonal instead**: Gkeyll writes +each mapped dimension as its own independent 1-D map (basis dimensionality +1, not ``m``), broadcast across the other velocity dimensions' cells -- +gyrokinetic velocity coordinates never couple. A real vel-space fixture +(``rt_gk_tcv_iwl_1x2v_p1-elc_mapc2p_vel.gkyl``) is laid out exactly this +way: its 4 components live on a 2-D (16, 8) grid, matching +``m * num_basis_1d == 2 * 2 == 4`` for 1-D serendipity p1, and it carries no +``basis_type``/``poly_order`` metadata of its own -- callers must supply it +(see ``load_distf``, which passes ``basis_type="serendipity", +poly_order=1``). +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") +pytestmark = needs_gkeyll + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GEN = os.path.join(DATA, "generated") +F_ELC = os.path.join(DATA, "rt_gk_tcv_iwl_1x2v_p1-elc_250.gkyl") +F_MAPC2P_VEL = os.path.join(DATA, "rt_gk_tcv_iwl_1x2v_p1-elc_mapc2p_vel.gkyl") + + +# --------------------------------------------------------------- test helpers +def _project_1d(fn, lower, upper, cells, basis_type, poly_order): + """Exact per-cell modal coefficients of ``fn(z)`` for a 1-D basis (see + ``tests/test_dg_map.py`` for the same helper at the engine level).""" + node_eta = gpython.basis.node_coords(basis_type, 1, poly_order)[:, 0] + n2m = gpython.basis.nodal_to_modal_matrix(basis_type, 1, poly_order) + dz = (upper - lower) / cells + centers = lower + (np.arange(cells) + 0.5) * dz + nodal_z = centers[:, None] + 0.5 * dz * node_eta[None, :] + return fn(nodal_z) @ n2m.T + + +def _project_2d(fn, lower, upper, cells, basis_type, poly_order): + """Exact per-cell modal coefficients of ``fn(z0, z1)`` for a 2-D basis.""" + node_eta = gpython.basis.node_coords(basis_type, 2, poly_order) + n2m = gpython.basis.nodal_to_modal_matrix(basis_type, 2, poly_order) + dz = [(upper[d] - lower[d]) / cells[d] for d in range(2)] + c0 = lower[0] + (np.arange(cells[0]) + 0.5) * dz[0] + c1 = lower[1] + (np.arange(cells[1]) + 0.5) * dz[1] + centers = np.stack(np.meshgrid(c0, c1, indexing="ij"), axis=-1) + node_phys = ( + centers[:, :, None, :] + + 0.5 * np.array(dz)[None, None, None, :] * node_eta[None, None, :, :]) + nodal_vals = fn(node_phys[..., 0], node_phys[..., 1]) + return np.einsum("ij,...j->...i", n2m, nodal_vals) + + +def _synthetic_map(coeffs, + lower, + upper, + cells, + *, + basis_type="serendipity", + poly_order=1, + value_form="modal"): + """A gkyl-backed mapping dataset holding ``coeffs`` directly -- no mapc2p + file needed, per the layer instructions. ``cells`` must be set in ``ctx`` + before ``push`` (``GDataState.set_grid`` needs it to know ``num_dims``, + and a flat ``GkylArray`` carries no cell layout of its own).""" + d = GDataState() + d.ctx.update(basis_type=basis_type, + poly_order=poly_order, + value_form=value_form, + cells=np.asarray(cells, dtype=np.int64)) + grid = [ + np.linspace(lower[i], upper[i], + int(cells[i]) + 1) for i in range(len(cells)) + ] + d.push(grid, gpython.GkylArray.from_numpy(coeffs)) + return d + + +def _numpy_target(grid, values): + """A NumPy-backed (field-domain) target dataset, built directly.""" + d = GDataState() + d.push(list(grid), values) + return d + + +# ----------------------------------------------------------------- identity +class TestIdentityMap: + + def test_1d_conf_identity_leaves_grid_unchanged(self): + lower, upper, cells = 0.0, 4.0, 4 + modal = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + mapping = _synthetic_map(modal, [lower], [upper], [cells]) + + target_axis = np.linspace(lower, upper, 17) # finer than the map's grid + target = _numpy_target([target_axis], np.zeros((16, 1))) + out = operations.map(target, mapping, space="conf") + + np.testing.assert_allclose(out.grid[0], target_axis, atol=1e-12) + assert out.ctx["grid_type"] == "mapped" + + def test_values_are_untouched(self): + lower, upper, cells = 0.0, 2.0, 2 + modal = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + mapping = _synthetic_map(modal, [lower], [upper], [cells]) + values = np.arange(8.0).reshape(4, 2) + target = _numpy_target([np.linspace(lower, upper, 5)], values) + out = operations.map(target, mapping, space="conf") + np.testing.assert_array_equal(out.values, values) + + def test_new_dataset_by_default_source_grid_untouched(self): + lower, upper, cells = 0.0, 2.0, 2 + modal = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + mapping = _synthetic_map(modal, [lower], [upper], [cells]) + target = _numpy_target([np.linspace(lower, upper, 5)], np.zeros((4, 1))) + out = operations.map(target, mapping, space="conf") + assert out is not target + assert "grid_type" not in target.ctx + + def test_inplace_mutates(self): + lower, upper, cells = 0.0, 2.0, 2 + modal = _project_1d(lambda z: z, lower, upper, cells, "serendipity", 1) + mapping = _synthetic_map(modal, [lower], [upper], [cells]) + target = _numpy_target([np.linspace(lower, upper, 5)], np.zeros((4, 1))) + out = operations.map(target, mapping, space="conf", inplace=True) + assert out is target + + +# ------------------------------------------------------------ conf, 2-D real +class TestConfMapRealFixture: + """The real generated ``2d_c2p_*`` fixtures for conf-space.""" + + def _mapped(self, mapfile): + # operations.map, not the fluent .map() -- api/gdata.py's fluent wiring for the + # new physics/map verbs is a different layer's job (out of this layer's + # scope; see the report). + data = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")).interpolate() + return operations.map(data, os.path.join(GEN, mapfile), space="conf") + + def test_grid_becomes_curvilinear_with_shape_of_the_axes_it_replaces(self): + before = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")).interpolate() + mapped = self._mapped("2d_c2p_stretch_ms_p1.gkyl") + expected_shape = (before.grid[0].shape[0], before.grid[1].shape[0]) + assert mapped.grid[0].shape == expected_shape + assert mapped.grid[1].shape == expected_shape + assert mapped.grid[0].ndim == 2 # curvilinear: full N-D nodal array + + def test_values_untouched_by_stretch_map(self): + before = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")).interpolate() + mapped = self._mapped("2d_c2p_stretch_ms_p1.gkyl") + np.testing.assert_array_equal(mapped.values, before.values) + + def test_rotation_is_non_separable(self): + """A rotation map produces coordinates that vary along both axes.""" + mapped = self._mapped("2d_c2p_rot45_ms_p1.gkyl") + assert np.std(mapped.grid[0], axis=1).max() > 1e-6 + + +# --------------------------------------------------------------------- vel +class TestVelMap: + + def test_1d_vel_deforms_only_the_trailing_axis(self): + """m=1: offset = num_dims - m puts the map on the last axis.""" + lower, upper, cells = -1.0, 1.0, 4 + scale = 2.0 + modal = _project_1d(lambda v: scale * v, lower, upper, cells, "serendipity", + 1) + mapping = _synthetic_map(modal, [lower], [upper], [cells]) + + x_edges = np.linspace(0.0, 1.0, 5) + v0_edges = np.linspace(0.0, 1.0, 5) + v1_edges = np.linspace(lower, upper, 9) + target = _numpy_target([x_edges, v0_edges, v1_edges], np.zeros( + (4, 4, 8, 1))) + out = operations.map(target, mapping, space="vel") + + np.testing.assert_allclose(out.grid[0], x_edges) # untouched + np.testing.assert_allclose(out.grid[1], v0_edges) # untouched + np.testing.assert_allclose(out.grid[2], scale * v1_edges, atol=1e-12) + + def test_2d_vel_is_separable_per_dimension(self): + """Gkeyll's real ``mapc2p_vel`` files (see the module docstring) store + each velocity dimension as its own independent 1-D map, broadcast + across the other dimensions' cells -- not a joint m-D curvilinear map + like a conf-space map. Both dimensions evaluate independently even + though m == 2, and the resulting grid arrays stay 1-D.""" + lower, upper, cells = [-1.0, -1.0], [1.0, 1.0], [4, 3] + m0 = _project_1d(lambda v: 2.0 * v, lower[0], upper[0], cells[0], + "serendipity", 1) + m1 = _project_1d(lambda v: 3.0 * v + 1.0, lower[1], upper[1], cells[1], + "serendipity", 1) + # Broadcast each 1-D map's coefficients across the other axis' cells, + # matching Gkeyll's on-disk mapc2p_vel layout. + coeffs0 = np.broadcast_to(m0[:, None, :], + (cells[0], cells[1], m0.shape[-1])) + coeffs1 = np.broadcast_to(m1[None, :, :], + (cells[0], cells[1], m1.shape[-1])) + mapping = _synthetic_map(np.concatenate([coeffs0, coeffs1], axis=-1), lower, + upper, cells) + + x_edges = np.linspace(0.0, 1.0, 3) + v0_edges = np.linspace(lower[0], upper[0], 9) + v1_edges = np.linspace(lower[1], upper[1], 7) + target = _numpy_target([x_edges, v0_edges, v1_edges], np.zeros( + (2, 8, 6, 1))) + out = operations.map(target, mapping, space="vel") + + assert out.grid[1].ndim == 1 # separable: stays 1-D unlike a conf map + assert out.grid[2].ndim == 1 + np.testing.assert_allclose(out.grid[1], 2.0 * v0_edges, atol=1e-12) + np.testing.assert_allclose(out.grid[2], 3.0 * v1_edges + 1.0, atol=1e-12) + np.testing.assert_allclose(out.grid[0], x_edges) # conf axis untouched + + +# --------------------------------------------------------------------- errors +class TestMapErrors: + + def test_rejects_modal_target(self): + target = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")) # not interpolated + mapping_path = os.path.join(GEN, "2d_c2p_stretch_ms_p1.gkyl") + with pytest.raises(ValueError, match=r"\.interpolate\(\)"): + operations.map(target, mapping_path, space="conf") + + def test_bad_space_raises(self): + target = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")).interpolate() + with pytest.raises(ValueError, match="'space'"): + operations.map(target, + os.path.join(GEN, "2d_c2p_stretch_ms_p1.gkyl"), + space="bogus") + + def test_map_too_large_for_dataset(self): + target = pg.load(os.path.join(GEN, "1d_ms_p1.gkyl")).interpolate() # 1-D + with pytest.raises(ValueError, match="does not fit"): + operations.map(target, + os.path.join(GEN, "2d_c2p_stretch_ms_p1.gkyl"), + space="conf") # a 2-D map does not fit 1-D data + + def test_num_comps_validation_error(self): + lower, upper, cells = 0.0, 1.0, 2 + bad = np.zeros((cells, 3)) # serendipity p1 1-D needs num_basis=2, not 3 + mapping = _synthetic_map(bad, [lower], [upper], [cells]) + target = _numpy_target([np.linspace(lower, upper, 5)], np.zeros((4, 1))) + with pytest.raises(ValueError, match="component"): + operations.map(target, mapping, space="conf") + + def test_missing_basis_metadata_raises(self): + d = GDataState() + d.ctx.update(cells=np.array([2])) + d.push([np.linspace(0.0, 1.0, 3)], + gpython.GkylArray.from_numpy(np.zeros((2, 2)))) + target = _numpy_target([np.linspace(0.0, 1.0, 5)], np.zeros((4, 1))) + with pytest.raises(ValueError, match="basis_type"): + operations.map(target, d, space="conf") + + def test_vel_map_real_fixture_fits_the_separable_algorithm(self): + """See the module docstring: this real fixture carries no basis + metadata of its own (callers must supply it at load time, as + ``load_distf`` does), and its 4 components match m * num_basis_1d + == 2 * 2 == 4 for 1-D serendipity p1 -- not m * num_basis_2d (== 2 * 4 + == 8), which is what the joint-curvilinear (conf-style) contract would + require.""" + mapping = pg.load(F_MAPC2P_VEL, basis_type="serendipity", poly_order=1) + assert mapping.ctx.get("basis_type") == "serendipity" + assert mapping.num_dims == 2 and mapping.num_comps == 4 + assert gpython.basis.num_basis("serendipity", 1, 1) == 2 + assert gpython.basis.num_basis("serendipity", 2, 1) == 4 + + target = pg.load(F_ELC).interpolate() + out = operations.map(target, mapping, space="vel") + assert out.grid[-1].ndim == 1 and out.grid[-2].ndim == 1 + assert np.all(np.isfinite(out.grid[-1])) + assert np.all(np.isfinite(out.grid[-2])) + + +# --------------------------------------- select on curvilinear (conf) grids +class TestSelectCurvilinearGuard: + + def _mapped(self): + data = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")).interpolate() + return operations.map(data, + os.path.join(GEN, "2d_c2p_rot45_ms_p1.gkyl"), + space="conf") + + def test_coordinate_selector_on_non_separable_axis_refuses(self): + """A rotation map genuinely couples both axes -- a bare coordinate + value has no single answer until the other axis is pinned first.""" + mapped = self._mapped() + with pytest.raises(ValueError, match="varies along another axis"): + mapped.select(z0=0.0) + + def test_slice_selector_on_non_separable_axis_refuses(self): + mapped = self._mapped() + with pytest.raises(ValueError, match="varies along another axis"): + mapped.select(z0="1:3") + + def test_coordinate_selector_works_once_the_other_axis_is_pinned(self): + """Selecting z1 by index first narrows its *values* to one cell (even + though the curvilinear grid array itself keeps 2 bounding edges); a + later, separate select() call's z0 coordinate curve is then read off + that specific (now unambiguous) cross-section -- resolved purely from + the dataset's own values shape, with no extra state to thread through.""" + mapped = self._mapped() + pinned = operations.select(mapped, z1=2) + assert pinned.values.shape[1] == 1 # z1 resolved to a single cell + assert pinned.grid[0].shape[1] == 2 # edges: 2 bound that one cell + out = operations.select(pinned, z0=pinned.grid[0][3, 0]) + assert out.values.shape[0] == 1 + + def test_integer_index_selector_still_works(self): + mapped = self._mapped() + out = mapped.select(z0=1) + assert out.values.shape[0] == 1 + # grid holds edges (2 bound one cell) even along a curvilinear axis + assert out.grid[0].shape[0] == 2 + + def test_coordinate_selector_works_on_a_separable_joint_map(self): + """A conf-space map stored jointly (m*num_basis components) but whose + physical coordinates each happen to depend on only one computational + axis -- e.g. Gkeyll's field-aligned ``mc2nu`` remap -- resolves a + coordinate value exactly, with no need to pin the other axis first.""" + lower, upper, cells = [0.0, 0.0], [1.0, 1.0], [2, 2] + m0 = _project_2d(lambda z0, z1: 2.0 * z0, lower, upper, cells, + "serendipity", 1) + m1 = _project_2d(lambda z0, z1: 3.0 * z1 + 1.0, lower, upper, cells, + "serendipity", 1) + mapping = _synthetic_map(np.concatenate([m0, m1], axis=-1), lower, upper, + cells) + target = _numpy_target([ + np.linspace(lower[0], upper[0], 5), + np.linspace(lower[1], upper[1], 5) + ], np.zeros((4, 4, 1))) + mapped = operations.map(target, mapping, space="conf") + assert mapped.grid[0].ndim == 2 # stored jointly, curvilinear shape + + out = operations.select(mapped, z0=1.0) # 2.0 * 0.5 == 1.0 + assert out.values.shape[0] == 1 + assert out.grid[0].shape == (2, 5) # only z0's own axis narrowed + assert out.grid[1].shape == (2, 5) # sibling kept in sync + + def test_selecting_one_axis_narrows_the_sibling_grid_too(self): + mapped = self._mapped() + out = operations.select(mapped, z0=1) + assert out.grid[0].shape[0] == 2 # z0's own axis: sliced 5 -> 2 (edges) + assert out.grid[1].shape[0] == 2 # z1's array shares the block shape + + def test_separable_1d_mapped_axis_keeps_coordinate_selection(self): + """A vel (m=1) mapped axis stays 1-D, so the ordinary coordinate-lookup + path (unaffected by the curvilinear guard) still applies.""" + lower, upper, cells = -1.0, 1.0, 4 + modal = _project_1d(lambda v: v, lower, upper, cells, "serendipity", 1) + mapping = _synthetic_map(modal, [lower], [upper], [cells]) + target = _numpy_target( + [np.linspace(0.0, 1.0, 5), + np.linspace(lower, upper, 9)], np.zeros((4, 8, 1))) + mapped = operations.map(target, mapping, space="vel") + assert mapped.grid[1].ndim == 1 + out = operations.select(mapped, z1=0.0) + assert out.values.shape[1] == 1 + + def test_2d_vel_map_keeps_ordinary_selection_behind_a_nonzero_offset(self): + """An m > 1 ``space="vel"`` map sits behind a nonzero ``offset`` + (``num_dims - m``), same as any vel map -- but because each mapped + dimension stays its own independent 1-D map (unlike a conf-space + curvilinear map), every mapped axis keeps ordinary coordinate-based + ``select``, with no ``mapped_axes`` offset translation needed.""" + lower, upper, cells = [-1.0, -1.0], [1.0, 1.0], [4, 3] + m0 = _project_1d(lambda v: 2.0 * v, lower[0], upper[0], cells[0], + "serendipity", 1) + m1 = _project_1d(lambda v: v, lower[1], upper[1], cells[1], "serendipity", + 1) + coeffs0 = np.broadcast_to(m0[:, None, :], + (cells[0], cells[1], m0.shape[-1])) + coeffs1 = np.broadcast_to(m1[None, :, :], + (cells[0], cells[1], m1.shape[-1])) + mapping = _synthetic_map(np.concatenate([coeffs0, coeffs1], axis=-1), lower, + upper, cells) + + x_edges = np.linspace(0.0, 1.0, 3) + v0_edges = np.linspace(lower[0], upper[0], 6) # non-square vs. v1 + v1_edges = np.linspace(lower[1], upper[1], 4) + target = _numpy_target([x_edges, v0_edges, v1_edges], + np.arange(2 * 5 * 3).reshape(2, 5, 3, + 1).astype(float)) + out = operations.map(target, mapping, space="vel") # offset = 3 - 2 = 1 + assert out.grid[1].ndim == 1 and out.grid[2].ndim == 1 + + sel2 = operations.select(out, z2=v1_edges[2]) + assert sel2.values.shape == (2, 5, 1, 1) + assert sel2.grid[2].shape == (2, ) # v1's own axis sliced 4 -> 2 + assert sel2.grid[1].shape == (6, ) # untouched by this call + + sel1 = operations.select(out, z1=v0_edges[1]) + assert sel1.values.shape == (2, 1, 3, 1) + assert sel1.grid[1].shape == (2, ) # v0's own axis sliced 6 -> 2 + assert sel1.grid[2].shape == (4, ) # untouched by this call diff --git a/tests/test_operations_sort.py b/tests/test_operations_sort.py new file mode 100644 index 00000000..c43bab5d --- /dev/null +++ b/tests/test_operations_sort.py @@ -0,0 +1,65 @@ +"""Tests for the ``sort`` verb and its ``numerics.natural_sort_key`` helper.""" + +from __future__ import annotations + +from postgkyl import numerics, operations +from postgkyl.gdatastate.gdatastate import GDataState + + +def _named(file_name): + d = GDataState() + d._file_name = file_name + return d + + +def test_natural_sort_key_orders_embedded_numbers_numerically(): + names = ["field_10.gkyl", "field_1.gkyl", "field_2.gkyl", "field_20.gkyl"] + assert sorted(names, key=numerics.natural_sort_key) == [ + "field_1.gkyl", "field_2.gkyl", "field_10.gkyl", "field_20.gkyl" + ] + + +def test_natural_sort_key_beats_plain_lexicographic_sort(): + # The whole point of natural sort: a plain string sort gets this wrong. + names = ["field_0.gkyl", "field_1.gkyl", "field_10.gkyl", "field_2.gkyl"] + assert sorted(names) != sorted(names, key=numerics.natural_sort_key) + assert sorted(names, key=numerics.natural_sort_key) == [ + "field_0.gkyl", "field_1.gkyl", "field_2.gkyl", "field_10.gkyl" + ] + + +def test_sort_reorders_datasets_by_filename(): + a = _named("field_10.gkyl") + b = _named("field_2.gkyl") + out = operations.sort(a, b) + assert [d.file_name for d in out] == ["field_2.gkyl", "field_10.gkyl"] + + +def test_sort_accepts_a_list_argument(): + frames = [ + _named("field_10.gkyl"), + _named("field_2.gkyl"), + _named("field_1.gkyl") + ] + out = operations.sort(frames) + assert [d.file_name + for d in out] == ["field_1.gkyl", "field_2.gkyl", "field_10.gkyl"] + + +def test_sort_reverse(): + a = _named("field_1.gkyl") + b = _named("field_2.gkyl") + out = operations.sort(a, b, reverse=True) + assert [d.file_name for d in out] == ["field_2.gkyl", "field_1.gkyl"] + + +def test_sort_does_not_mutate_or_copy_datasets(): + a = _named("field_2.gkyl") + b = _named("field_1.gkyl") + out = operations.sort(a, b) + assert out == [b, a] + assert out[0] is b and out[1] is a + + +def test_sort_empty_returns_empty(): + assert operations.sort() == [] diff --git a/tests/test_packaging.py b/tests/test_packaging.py new file mode 100644 index 00000000..97be53cf --- /dev/null +++ b/tests/test_packaging.py @@ -0,0 +1,44 @@ +"""Regression tests for the custom native packaging commands.""" + +import runpy +from pathlib import Path + +import setuptools +from setuptools.dist import Distribution + +ROOT_DIR = Path(__file__).parents[1] + + +def _build_py_command(monkeypatch, tmp_path, *, editable): + monkeypatch.setattr(setuptools, "setup", lambda **kwargs: None) + namespace = runpy.run_path(ROOT_DIR / "setup.py") + command_type = namespace["BuildPyWithGkeyll"] + native_library = tmp_path / "libg0core.so" + native_library.write_bytes(b"native library") + monkeypatch.setitem(command_type.run.__globals__, "_build_gkeyll", + lambda: True) + monkeypatch.setitem(command_type.run.__globals__, "BUNDLED_LIB", + native_library) + + command = command_type(Distribution()) + command.ensure_finalized() + command.build_lib = str(tmp_path / "missing-build") + command.editable_mode = editable + return command + + +def test_editable_build_uses_native_artifacts_in_source(monkeypatch, tmp_path): + command = _build_py_command(monkeypatch, tmp_path, editable=True) + + command.run() + + assert not Path(command.build_lib).exists() + + +def test_wheel_build_creates_native_library_destination(monkeypatch, tmp_path): + command = _build_py_command(monkeypatch, tmp_path, editable=False) + + command.run() + + bundled = Path(command.build_lib) / "postgkyl/gpython/libg0core.so" + assert bundled.read_bytes() == b"native library" diff --git a/tests/test_plot.py b/tests/test_plot.py deleted file mode 100644 index e38d41ea..00000000 --- a/tests/test_plot.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Postgkyl module for testing the plotting function""" -import os -import matplotlib as mpl -import numpy as np - -import postgkyl as pg - -class TestPlot: - """Test Postgkyl plot function. - - Currently, this tests if plots look OK only to some extend (by checking plotted - values) and mostly tests if plots are created at all. Testing images themselves is - complicated and differs based on system and/or backend used. - """ - dir_path = f"{os.path.dirname(__file__)}/test_data" - - def test_plot_pcolormesh(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - img = pg.output.plot(data) - assert isinstance(img, mpl.collections.QuadMesh) - mpl.pyplot.close("all") - - def test_plot_contour(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - img = pg.output.plot(data, contour=True) - assert isinstance(img, mpl.contour.QuadContourSet) - mpl.pyplot.close("all") - - def test_plot_contour_options(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - img = pg.output.plot(data, contour=True, cnlevels=5, cont_label=True) - assert isinstance(img, mpl.contour.QuadContourSet) - mpl.pyplot.close("all") - - def test_plot_line(self): - data = pg.GData(f"{self.dir_path:s}/twostream-field-energy.gkyl") - img = pg.output.plot(data) - assert isinstance(img[0], mpl.lines.Line2D) - mpl.pyplot.close("all") - - pg.data.select(data, comp=0, overwrite=True) - img = pg.output.plot(data) - x_plot, y_plot = img[0].get_xydata().T - np.testing.assert_array_almost_equal(data.get_grid()[0], x_plot) - np.testing.assert_array_almost_equal(data.get_values()[...,0], y_plot) - mpl.pyplot.close("all") \ No newline at end of file diff --git a/tests/test_postgkyl.py b/tests/test_postgkyl.py new file mode 100644 index 00000000..03b62702 --- /dev/null +++ b/tests/test_postgkyl.py @@ -0,0 +1,687 @@ +"""Smoke tests + architecture contract for the postgkyl library. + +Run: PYTHONPATH=src pytest tests/test_postgkyl.py -v +""" + +import ast +import collections +import os +import sys +from pathlib import Path + +import numpy as np +import pytest + +# Make src/ importable without an install. +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SRC = os.path.join(ROOT, "src") +sys.path.insert(0, SRC) # dedup harmless across the shared test session + +import matplotlib + +matplotlib.use("Agg") + +import postgkyl as pg # noqa: E402 + +DATA = os.path.join(ROOT, "tests", "test_data") +F1 = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F2D = os.path.join(DATA, "generated", "2d_ms_p1.gkyl") +F_GKHYBRID = os.path.join(DATA, "rt_gk_tcv_iwl_1x2v_p1-elc_250.gkyl") + + +def test_load_metadata(): + d = pg.load(F1) + assert d.num_dims == 1 + assert d.ctx["basis_type"] == "serendipity" + assert d.ctx["poly_order"] == 1 + assert not d.is_interpolated # raw modal data + + +def test_golden_script_1d(): + g = pg.load(F1).interpolate().select(comp=0) + assert g.is_interpolated + assert g.num_comps == 1 + assert g.num_dims == 1 + assert g.values.shape[0] == 48 # 24 cells * (p+1=2) interpolation points + assert type(g).__name__ == "GData" # subclass propagated through verbs + fig = g.plot(no_show=True) + assert fig is not None + + +def test_golden_script_2d(): + g = pg.load(F2D).interpolate().select(comp=0) + assert g.num_dims == 2 + assert g.values.shape == (16, 16, 1) + assert g.plot(no_show=True) is not None + + +def test_plot_has_one_canonical_callable(): + from postgkyl import operations, render + from postgkyl.gdata import verbs + + assert pg.plot is render.plot + assert pg.plot is operations.plot + assert pg.plot is pg.GData.plot + assert pg.plot is pg.GDataGroup.plot + assert pg.plot is verbs.plot + + +def test_plotly_has_one_canonical_callable(): + from postgkyl import operations, render + + assert pg.plotly is render.plotly + assert pg.plotly is operations.plotly + assert pg.plotly is pg.GData.plotly + + +def test_pyvista_has_one_canonical_callable(): + from postgkyl import operations, render + + assert pg.pyvista is render.pyvista + assert pg.pyvista is operations.pyvista + assert pg.pyvista is pg.GData.pyvista + + +def test_arithmetic_and_ufunc(): + a = pg.load(F1).interpolate().select(comp=0) + b = pg.load(F1).interpolate().select(comp=0) + assert isinstance(a + b, pg.GData) + assert isinstance(a * 2.0, pg.GData) + assert isinstance(2.0 * a, pg.GData) # reflected + mag = np.sqrt(a**2 + b**2) # ufunc keeps it a GData + assert isinstance(mag, pg.GData) + assert np.allclose(mag.values, np.sqrt(a.values**2 + b.values**2)) + assert np.asarray(a).shape == a.values.shape # __array__ + + +def test_capability_guardrails_on_modal_data(): + """Modal data supports the Gkeyll verbs; everything NumPy-shaped refuses.""" + a = pg.load(F1) + with pytest.raises(ValueError): + np.sqrt(a) # general ufunc: no modal meaning + with pytest.raises(ValueError): + np.asarray(a) # coefficients are not point values + with pytest.raises(ValueError): + a.select(comp=0) # slicing would mix basis functions + with pytest.raises(ValueError): + _ = a + a.interpolate() # mixed modal + field domains + + +# -------------------------------------------------------------------------- +# The modal domain: DG operations running inside Gkeyll (REFACTOR_GKEYLL_FFI.md) +# -------------------------------------------------------------------------- +from postgkyl import gpython # noqa: E402 + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +@needs_gkeyll +def test_load_lands_in_the_modal_domain(): + d = pg.load(F1) + assert d.backend == "gkyl" # native gkyl_array storage + assert d.native is not None + assert d.values.shape == (24, 6) # read-only view for inspection + assert not d.values.flags.writeable + g = d.interpolate() # the one-way bridge + assert g.backend == "numpy" # ...to a by-value NumPy array + assert g.values.flags.writeable + + +@needs_gkeyll +def test_shim_handshake(): + """The compiled gpython shim pairs with this postgkyl (GKEYLL_C_SHIM.md). + + There are no struct layouts to guard anymore -- the C compiler checked the + whole contract when gpython.c built. What remains testable at runtime is the + version handshake plus a behavioral probe through the shim.""" + g0 = gpython.require() + assert g0.api_version() == g0.GPYTHON_API_VERSION + b = gpython.basis.get_basis("serendipity", 2, 1) + assert (b.ndim, b.poly_order, b.num_basis) == (2, 1, 4) + assert b.id == "serendipity" + + +@needs_gkeyll +def test_gkhybrid_basis_loads_and_interpolates(): + """A real 1x2v gyrokinetic distribution file (gkhybrid basis) round-trips + through the modal -> field bridge, exactly like a serendipity/tensor file.""" + d = pg.load(F_GKHYBRID) + assert d.ctx["basis_type"] == "gkhybrid" + assert d.ctx["poly_order"] == 1 + assert d.num_dims == 3 # 1x2v + assert d.values.shape[-1] == 12 # gkhybrid 1x2v num_basis + g = d.interpolate() + assert g.backend == "numpy" + assert g.values.shape == (64, 32, 16, 1) # (p+1=2) interpolation points/cell + + +@needs_gkeyll +def test_interpolation_matrix_matches_analytic_basis(): + """Matrices built from Gkeyll's eval() match the normalized Legendre basis.""" + m = gpython.basis.interpolation_matrix("serendipity", 1, 1, + 2) # points z = -+1/2 + expect = np.array([[1 / np.sqrt(2), -np.sqrt(3.0 / 2.0) / 2], + [1 / np.sqrt(2), +np.sqrt(3.0 / 2.0) / 2]]) + assert np.allclose(m, expect) + m2 = gpython.basis.interpolation_matrix("serendipity", 1, 2, + 3) # p2, points -+2/3, 0 + z = np.array([-2.0 / 3.0, 0.0, 2.0 / 3.0]) + assert np.allclose(m2[:, 2], 2.371708245126285 * z**2 - 0.7905694150420951) + + +@needs_gkeyll +def test_weak_algebra_identities(): + """div(mul(a, b), b) == a -- Gkeyll's weak kernels are exact inverses.""" + a, b = pg.load(F1), pg.load(F1) + back = (a * b / b).interpolate().values + ref = a.interpolate().values + for f in (0, 2): # density and T; field 1 (u_par) is identically ~0 -> 0/0 + scale = np.abs(ref[..., f]).max() + assert np.abs(back[..., f] - ref[..., f]).max() / scale < 1e-12 + + +@needs_gkeyll +def test_modal_linear_ops_commute_with_interpolate(): + """interpolate is linear: modal +,-,scalar* agree with their NumPy counterparts.""" + a, b = pg.load(F1), pg.load(F1) + assert np.allclose((a + b).interpolate().values, + a.interpolate().values + b.interpolate().values) + assert np.allclose((a - b).interpolate().values, 0.0) + assert np.allclose((2.5 * a).interpolate().values, + 2.5 * a.interpolate().values) + assert np.allclose((-a).interpolate().values, -(a.interpolate().values)) + assert np.allclose((a**2).interpolate().values, (a * a).interpolate().values) + shifted = (a + 1.0e18).interpolate().values - a.interpolate().values + assert np.allclose(shifted, 1.0e18, rtol=1e-6) + + +def _make_modal(grid, cells, basis_type, poly_order, coeffs): + """A bare modal GData, built in-memory rather than from a file -- for + exercising the conf x phase cross-multiply path with grids we control.""" + d = pg.GData() + d.ctx.update(basis_type=basis_type, + poly_order=poly_order, + value_form="modal", + cells=np.array(cells)) + d.push(grid, gpython.array.GkylArray.from_numpy(coeffs)) + return d + + +@needs_gkeyll +def test_conf_phase_mul_is_automatic_and_commutative(): + """``conf * phase`` and ``phase * conf`` both dispatch to the cross-basis + gkyl_dg_mul_conf_phase_op_range path with no separate method needed -- the + API picks the lower-dimensional operand as the conf side automatically. + Multiplying by a spatially-uniform conf field of true value 1 is an exact + identity on the phase side (no weak-projection truncation), so this is a + correctness check, not just a "did it run" smoke test.""" + conf_edges = [np.linspace(0.0, 1.0, 4)] # 3 cells + phase_edges = [np.linspace(0.0, 1.0, 4), np.linspace(-1.0, 1.0, 5)] # 3x4 + + cbasis = gpython.basis.get_basis("serendipity", 1, 1) + pbasis = gpython.basis.get_basis("hybrid", 2, 1) + cop = np.zeros((3, cbasis.num_basis)) + cop[:, 0] = np.sqrt(2.0) # value 1 + rng = np.random.default_rng(11) + pop = rng.normal(size=(12, pbasis.num_basis)) + + conf = _make_modal(conf_edges, [3], "serendipity", 1, cop) + phase = _make_modal(phase_edges, [3, 4], "hybrid", 1, pop) + + out1 = conf * phase + out2 = phase * conf + assert isinstance(out1, pg.GData) and out1.num_dims == 2 + np.testing.assert_allclose(out1.values.reshape(12, 6), pop) + np.testing.assert_allclose(out2.values.reshape(12, 6), pop) + + +@needs_gkeyll +def test_conf_phase_mul_rejects_non_mul_ops_and_mismatched_grids(): + conf = _make_modal([np.linspace(0.0, 1.0, 4)], [3], "serendipity", 1, + np.zeros((3, 2))) + phase = _make_modal([np.linspace(0.0, 1.0, 4), + np.linspace(-1.0, 1.0, 5)], [3, 4], "hybrid", 1, + np.zeros((12, 6))) + with pytest.raises(ValueError, match="only '\\*' is defined"): + conf / phase + with pytest.raises(ValueError, match="only '\\*' is defined"): + conf + phase + + mismatched = _make_modal( + [np.linspace(0.0, 2.0, 4), + np.linspace(-1.0, 1.0, 5)], [3, 4], "hybrid", 1, np.zeros((12, 6))) + with pytest.raises(ValueError, match="not the same simulation"): + conf * mismatched + + +def _relerr(x, y): + x, y = np.asarray(x, float), np.asarray(y, float) + return np.abs(x - y).max() / np.abs(y).max() + + +@needs_gkeyll +def test_value_form_round_trips(): + """modal <-> nodal is exact; modal <-> quad is exact for num_quad >= p+1.""" + a = pg.load(F1) + n = a.to_nodal() + assert n.ctx["value_form"] == "nodal" + assert n.backend == "gkyl" # never leaves the native domain + assert _relerr(n.to_modal().values, a.values) < 1e-14 + q = a.to_quad() + assert (q.ctx["value_form"], q.ctx["num_quad"]) == ("quad", 2) + assert _relerr(q.to_modal().values, a.values) < 1e-14 + # nodal -> quad composes through modal + assert _relerr(n.to_quad().to_modal().values, a.values) < 1e-14 + # nodal values are the field evaluated at the basis node_list points + m2n = gpython.basis.modal_to_nodal_matrix("serendipity", 1, 1) + manual = np.einsum("pk,cfk->cfp", m2n, + np.asarray(a.values).reshape(24, 3, 2)).reshape(24, 6) + assert np.allclose(n.values, manual) + + +@needs_gkeyll +def test_apply_pointwise_via_quadrature(): + """.apply(fn): modal -> quad -> fn -> modal, exact where quadrature is.""" + a = pg.load(F1) + assert _relerr(a.apply(lambda v: v).values, a.values) < 1e-13 + # p=1: p+1 Gauss points integrate the square exactly -> matches the weak kernel + assert _relerr(a.apply(np.square).values, (a * a).values) < 1e-13 + chained = a.apply(np.abs).apply(np.sqrt) # stays modal + gkyl-native + assert chained.backend == "gkyl" + assert chained.ctx.get("value_form", "modal") == "modal" + with pytest.raises(ValueError): + a.apply(lambda v: v.sum(axis=-1)) # fn must act pointwise + + +@needs_gkeyll +def test_conversions_are_always_explicit(): + """No implicit value_form change, ever (REFACTOR_GKEYLL_FFI.md §3b).""" + a = pg.load(F1) + n, q = a.to_nodal(), a.to_quad() + with pytest.raises(ValueError): + _ = a + n # mixed value_forms + with pytest.raises(ValueError): + _ = np.add(n, q) # mixed reps through a ufunc + with pytest.raises(ValueError): + q.interpolate() # interp needs modal + assert n.integrate() is not None # point values integrate in-place in form + with pytest.raises(ValueError): + np.sqrt(a) # ufuncs have no modal meaning + with pytest.raises(ValueError): + np.asarray(a) # coefficients are not values + with pytest.raises(ValueError): + a.plot(no_show=True) # coefficients are not plottable + + +@needs_gkeyll +def test_pointwise_numpy_on_point_values(): + """NumPy math is exact on nodal/quad data and stays native, in-value_form.""" + a = pg.load(F1) + n, q = a.to_nodal(), a.to_quad() + s = np.sqrt(np.abs(n)) # ufunc on nodal + assert (s.backend, s.ctx["value_form"]) == ("gkyl", "nodal") + assert np.allclose(s.values, np.sqrt(np.abs(np.asarray(n.values)))) + assert np.allclose((n**2).values, np.asarray(n.values)**2) + assert np.allclose((q * q).values, np.asarray(q.values)**2) + # pointwise-at-quad then one projection == the weak kernel (p1 exactness) + assert _relerr((q * q).to_modal().values, (a * a).values) < 1e-13 + # chain at the points, project once -- identical to the one-shot .apply() + fn = lambda v: np.sqrt(np.abs(v)) + assert _relerr(np.sqrt(np.abs(q)).to_modal().values, + a.apply(fn).values) < 1e-15 + assert np.asarray(n).shape == (24, 6) # __array__ allowed on points + + +@needs_gkeyll +def test_plot_point_values_directly(): + """Nodal/quad datasets plot at their true point locations.""" + a = pg.load(F1) + assert a.to_nodal().plot(no_show=True) is not None + assert a.to_quad().plot(no_show=True) is not None + b = pg.load(F2D) + assert b.to_quad().plot(no_show=True) is not None + assert b.to_nodal().plot(no_show=True) is not None # p1 corners: tensor set + p2 = pg.load(os.path.join(DATA, "generated", "2d_ms_p2.gkyl")) + with pytest.raises(ValueError): + p2.to_nodal().plot(no_show=True) # non-tensor node set -> to_quad + + +@needs_gkeyll +def test_linear_ops_valid_in_any_value_form(): + """+ - and scalar ops act pointwise in nodal/quad and agree with modal.""" + a = pg.load(F1) + n = a.to_nodal() + assert _relerr((2 * n - n + n).to_modal().values, (2 * a).values) < 1e-13 + assert _relerr((n + 5.0e17).to_modal().values, (a + 5.0e17).values) < 1e-13 + + +@needs_gkeyll +def test_values_view_pins_native_memory(): + """Regression: `dataset.values` on a temporary must stay valid after GC.""" + import gc + a = pg.load(F1) + expected = a.values.copy() + v = pg.load(F1).values # dataset is garbage immediately + got = (2 * pg.load(F1).to_nodal()).to_modal().values # temporaries galore + gc.collect() + assert np.array_equal(v, expected) + assert _relerr(got, 2 * expected) < 1e-13 + + +@needs_gkeyll +def test_integrate_via_gkeyll(): + """pg-level integrate == the coefficient-space formula (exact for DG).""" + a = pg.load(F1) + result = a.integrate() + v = a.values # (cells, nfields*num_basis) view + dx = float((a.bounds[1][0] - a.bounds[0][0]) / a.num_cells[0]) + nb = 2 # serendipity 1D p1 + manual = np.array([ + v[:, f * nb].sum() * dx / np.sqrt(2.0) for f in range(v.shape[-1] // nb) + ]) + assert np.allclose(result, manual) + assert np.all(a.integrate(op="abs") >= np.abs(result) * (1 - 1e-12)) + point_result = a.interpolate().integrate() + assert np.allclose(point_result, result) + + +def test_write_roundtrip(tmp_path): + a = pg.load(F1).interpolate().select(comp=0) + out = a.save(str(tmp_path / "rt.gkyl")) + back = pg.load(out) + assert np.allclose(back.values, a.values) + + +def test_info_returns_string(capsys): + s = pg.load(F1).info() + assert "Number of components" in s + + +def test_cli_chained(tmp_path): + """The chained CLI: bare filename -> load, interp, sel, plot --saveas.""" + from click.testing import CliRunner + from postgkyl.cli.app import cli + + out = tmp_path / "cli.png" + result = CliRunner().invoke(cli, [ + F1, "interp", "sel", "--comp", "0", "plot", "--no_show", "--saveas", + str(out) + ]) + assert result.exit_code == 0, result.output + assert out.exists() + + +def test_cli_abbreviation_and_info(): + """`interp`/`sel` resolve by unique-prefix abbreviation.""" + from click.testing import CliRunner + from postgkyl.cli.app import cli + + result = CliRunner().invoke(cli, [F1, "interp", "sel", "--comp", "0", "info"]) + assert result.exit_code == 0, result.output + assert "interpolated" in result.output + + +# -------------------------------------------------------------------------- +# Architecture contract: the layering is a strict, cycle-free DAG. +# -------------------------------------------------------------------------- +_ALLOWED = { + "cli_spec": set(), # frozen CLI metadata; dependency-free leaf + "gpython": set(), # the foreign floor (only ctypes owner) + "numerics": set(), + "dg": {"gpython"}, # interpolation bridge + modal ops -> kernels + "io": {"gpython", "numerics", + "cli_spec"}, # C-native reader -> gkyl_array_rio; + # writer reuses the pure-math leaf + # (nodal_to_cell_centered_grid for the + # vtk writer) instead of duplicating + # it -- numerics has 0 internal imports, + # so this cannot create a cycle (layer 04-io) + "gdatastate": {"io", "gpython", "dg"}, # state plus the shared native + # point-value materialization bridge + "render": {"gdatastate", "numerics", "cli_spec"}, + "operations": {"gdatastate", "dg", "io", "numerics", "render", + "cli_spec"}, # data transformations: + # print.py inspects stored values/grids through gdatastate only + # the physics verbs (moments/agyro/ + # current/energetics/rotate/ + # transform_frame/laguerre) moved up + # into diagnostics, folded with the + # models/ array math they delegated to; + # flat modules are equation-blind core + # verbs; domain subpackages (currently + # gyrokinetics) own transformations that + # need domain geometry without deriving + # a physical conclusion + "diagnostics": { + "gdatastate", "operations", "numerics", "gdata", "render", "io", + "cli_spec" + }, # added by + # 10-diagnostics.md: equation- + # specific compositions grouped under + # gk/vm/pkpm/mom; + # their modules wrap core + # verbs and state -- none of gdatastate/operations/ + # numerics imports upward, so this + # cannot create a cycle; "gdata" added by + # 12-diagnostics-loaders.md: the + # gk/pkpm loaders build on + # pg.load/GData (modal arithmetic, + # .interpolate()) to read simulation output + # -- gdata imports only gdatastate/operations/io, none + # of which import diagnostics, so this + # still cannot create a cycle; "render" + # pre-authorized by 13-diagnostics- + # programs.md for future program-scale + # diagnostics that may want render's + # generic plot() -- as of this layer's + # landing, none of the six program + # modules (energy_balance, particle_ + # balance, nodes, trajectory, enstrophy, + # ke_dke) actually import it, each + # building its own bespoke figure + # directly with matplotlib instead; + # render imports only gdatastate/numerics, + # neither of which imports diagnostics, + # so this cannot create a cycle whether + # or not the edge is ever exercised + "gdata": {"gdatastate", "operations", "io", "cli_spec"}, + "": { + "gdata", "operations", "render", "io", "gdatastate", "diagnostics", + "gpython", "_version", "cli_spec" + }, # facade: + # pure re-export of public names; + # "gdatastate" is group_blocks, the + # multiblock-family partition, which + # lives beside flatten_datasets in the + # container layer that owns collections + # of datasets; + # "diagnostics" added by + # 12-diagnostics-loaders.md, which + # explicitly authorizes the facade + # re-exporting the gk namespace for + # pg.gk.load_quantity(...); + # "_version" is __init__.py's own import + # of _version.py's version_report (`pgkyl + # --version`'s commit/build-info report), + # re-exported like any other facade name; + # "gpython" is _version.py's own edge (it + # reads gpython.available()/build_info()) + # -- both source files sit in the same "" + # layer, so both edges are checked here + "cli": {"", "cli_spec"}, # top surface: facade + frozen metadata +} +_LAYERS = set(_ALLOWED) + + +def _layer(path, pkg_root): + parts = os.path.relpath(path, pkg_root).split(os.sep) + if len(parts) > 1: + return parts[0] + module = os.path.splitext(parts[0])[0] + return module if module in _LAYERS else "" + + +def _import_targets(node): + if isinstance(node, ast.Import): + for n in node.names: + if n.name == "postgkyl" or n.name.startswith("postgkyl."): + t = n.name.split(".") + yield t[1] if len(t) > 1 else "" + elif isinstance(node, ast.ImportFrom): + if node.level: + return + mod = node.module or "" + if mod == "postgkyl": + for n in node.names: + yield n.name if n.name in _LAYERS else "" + elif mod.startswith("postgkyl."): + yield mod.split(".")[1] + + +def _build_edges(pkg_root=None): + pkg_root = pkg_root or os.path.join(SRC, "postgkyl") + edges = collections.defaultdict(set) + violations = [] + for dp, _, files in os.walk(pkg_root): + for f in files: + if not f.endswith(".py"): + continue + p = os.path.join(dp, f) + src = _layer(p, pkg_root) + for node in ast.walk(ast.parse(Path(p).read_text(encoding="utf-8"), p)): + for tgt in _import_targets(node): + if tgt == src: + continue + edges[src].add(tgt) + if tgt not in _ALLOWED.get(src, set()): + violations.append( + f"{os.path.relpath(p, pkg_root)} [{src or 'facade'}] -> [{tgt or 'facade'}]" + ) + return edges, violations + + +def test_facade_is_pure_reexport(): + """__init__.py must define no functions/classes -- only re-export names.""" + facade = os.path.join(SRC, "postgkyl", "__init__.py") + tree = ast.parse(Path(facade).read_text(encoding="utf-8"), facade) + defs = [ + n.name for n in tree.body + if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) + ] + assert not defs, f"facade should be pure re-export, but defines: {defs}" + + +def test_import_contract_no_violations(): + _, violations = _build_edges() + assert not violations, "layer contract violations:\n" + "\n".join(violations) + + +def _foreign_floor_offenders(pkg_root): + offenders = [] + for dp, _, files in os.walk(pkg_root): + for f in files: + if not f.endswith(".py"): + continue + p = os.path.join(dp, f) + in_gpython = _layer(p, pkg_root) == "gpython" + for node in ast.walk(ast.parse(Path(p).read_text(encoding="utf-8"), p)): + names = [] + if isinstance(node, ast.Import): + names = [n.name for n in node.names] + elif isinstance(node, ast.ImportFrom): + names = [node.module or "" + ] + ([n.name for n in node.names] + if node.level or "." in (node.module or "") or + (node.module or "") == "postgkyl" else []) + for name in names: + root = name.split(".")[0] + if root == "ctypes": + offenders.append(f"{os.path.relpath(p, pkg_root)}: ctypes") + if ("_gpython" in name.split(".") + or name == "_gpython") and not in_gpython: + offenders.append(f"{os.path.relpath(p, pkg_root)}: _gpython") + return offenders + + +def test_foreign_floor_confined_to_gpython(): + """The foreign world is the compiled ``_gpython`` extension, importable only + under gpython/ -- and ctypes appears nowhere at all: the C contract is enforced + by the compiler when the gpython shim builds, never re-declared in Python + (GKEYLL_C_SHIM.md).""" + pkg_root = os.path.join(SRC, "postgkyl") + offenders = _foreign_floor_offenders(pkg_root) + assert not offenders, f"foreign floor leaked above gpython/: {offenders}" + + +def _find_cycles(edges): + color = collections.defaultdict(int) + cycles = [] + + def dfs(u, stack): + color[u] = 1 + for w in edges.get(u, ()): + if color[w] == 1: + cycles.append(stack + [w]) + elif color[w] == 0: + dfs(w, stack + [w]) + color[u] = 2 + + for n in list(edges): + if color[n] == 0: + dfs(n, [n]) + return cycles + + +def test_import_graph_is_acyclic(): + edges, _ = _build_edges() + cycles = _find_cycles(edges) + assert not cycles, f"import cycle(s): {cycles}" + + +# -------------------------------------------------------------------------- +# The self-checks above only ever see a *compliant* tree in this repo (that +# is the point). These drive their violation/cycle/offender branches +# directly, against a small throwaway fake package tree, without touching +# the real source. +# -------------------------------------------------------------------------- +def _write_module(pkg_root, layer, name, body): + d = os.path.join(pkg_root, layer) if layer else pkg_root + os.makedirs(d, exist_ok=True) + with open(os.path.join(d, name), "w") as fh: + fh.write(body) + + +def test_build_edges_flags_a_disallowed_import(tmp_path): + pkg_root = str(tmp_path / "postgkyl") + _write_module(pkg_root, "badlayer", "mod.py", "import postgkyl.operations\n") + _, violations = _build_edges(pkg_root) + assert any("badlayer" in v and "operations" in v for v in violations) + + +def test_build_edges_classifies_a_flat_leaf_module(tmp_path): + pkg_root = str(tmp_path / "postgkyl") + _write_module(pkg_root, "", "cli_spec.py", "import postgkyl.operations\n") + _, violations = _build_edges(pkg_root) + assert any("cli_spec.py [cli_spec] -> [operations]" in v for v in violations) + + +def test_import_graph_detects_a_real_cycle(tmp_path): + pkg_root = str(tmp_path / "postgkyl") + _write_module(pkg_root, "layer_a", "mod.py", "import postgkyl.layer_b\n") + _write_module(pkg_root, "layer_b", "mod.py", "import postgkyl.layer_a\n") + edges, _ = _build_edges(pkg_root) + cycles = _find_cycles(edges) + assert cycles, "expected the fake layer_a <-> layer_b cycle to be detected" + + +def test_foreign_floor_offenders_flags_ctypes_and_gpython_outside_gpython( + tmp_path): + pkg_root = str(tmp_path / "postgkyl") + _write_module(pkg_root, "badlayer", "uses_ctypes.py", "import ctypes\n") + _write_module(pkg_root, "badlayer", "uses_gpython.py", + "from postgkyl.gpython import _gpython\n") + offenders = _foreign_floor_offenders(pkg_root) + assert any(o.endswith(": ctypes") for o in offenders) + assert any(o.endswith(": _gpython") for o in offenders) diff --git a/tests/test_render_animate.py b/tests/test_render_animate.py new file mode 100644 index 00000000..0309cec5 --- /dev/null +++ b/tests/test_render_animate.py @@ -0,0 +1,423 @@ +"""Tests for postgkyl.render.animate -- FuncAnimation / saved frames / movie +compile. + +Builds frames directly as ``GDataState`` (no shim dependency needed for the +render-layer tests. ``ffmpeg``-dependent tests are skipped cleanly when it is +not on ``PATH``. +""" + +from __future__ import annotations + +import os +from importlib import import_module + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.render import _ffmpeg + +anim_mod = import_module("postgkyl.render.animate") + +needs_ffmpeg = pytest.mark.skipif( + _ffmpeg.resolve_ffmpeg() is None, + reason="ffmpeg not found on PATH or via imageio-ffmpeg") +external_tool = pytest.mark.external_tool +slow = pytest.mark.slow + + +def _line_frame(offset: float) -> GDataState: + d = GDataState() + d.ctx["frame"] = int(offset) + d.ctx["time"] = float(offset) * 0.1 + d.push([np.linspace(0.0, 1.0, 9)], (np.arange(8, dtype=float) + offset)[:, + None]) + return d + + +def _three_frames() -> list[GDataState]: + return [_line_frame(0.0), _line_frame(1.0), _line_frame(2.0)] + + +@pytest.fixture(autouse=True) +def _close_figs(): + plt.close("all") + yield + plt.close("all") + + +# -------------------------------------------------------------------------- +# frame normalization +# -------------------------------------------------------------------------- + + +class TestNormalizeFrames: + + def test_bare_datasets_become_single_dataset_frames(self): + frames = anim_mod._normalize_frames(_three_frames()) + assert len(frames) == 3 + assert all(len(f) == 1 for f in frames) + + def test_grouped_frames_kept_as_lists(self): + grouped = [[_line_frame(0.0), _line_frame(0.5)], [_line_frame(1.0)]] + frames = anim_mod._normalize_frames(grouped) + assert len(frames) == 2 + assert len(frames[0]) == 2 + assert len(frames[1]) == 1 + + def test_empty_input_raises(self): + with pytest.raises(ValueError, match="no datasets"): + anim_mod._normalize_frames([]) + + +# -------------------------------------------------------------------------- +# fixed value range +# -------------------------------------------------------------------------- + + +class TestFrameValueRange: + + def test_spans_every_frame(self): + frames = anim_mod._normalize_frames(_three_frames()) + vmin, vmax = anim_mod._frame_value_range(frames) + assert vmin == 0.0 + assert vmax == 9.0 # last frame: arange(8) + 2.0 -> max 9.0 + + def test_cutoff_clips_the_range(self): + frames = anim_mod._normalize_frames(_three_frames()) + vmin_full, vmax_full = anim_mod._frame_value_range(frames) + vmin_cut, vmax_cut = anim_mod._frame_value_range(frames, cutoff=0.5) + assert vmin_cut >= vmin_full + assert vmax_cut <= vmax_full + + def test_scale_is_applied_before_taking_extrema(self): + # A fixed range computed on unscaled values would not match what + # matplotlib.plot actually draws once yscale/zscale is applied. + frames = anim_mod._normalize_frames(_three_frames()) + vmin, vmax = anim_mod._frame_value_range(frames, yscale=2.0) + assert vmin == 0.0 + assert vmax == 18.0 # last frame: (arange(8) + 2.0).max() * 2.0 + + +# -------------------------------------------------------------------------- +# live FuncAnimation path +# -------------------------------------------------------------------------- + + +class TestLiveAnimation: + + def test_returns_funcanimation_with_correct_frame_count(self): + from matplotlib.animation import FuncAnimation + anim = anim_mod.animate(_three_frames(), no_show=True) + assert isinstance(anim, FuncAnimation) + assert anim._save_count == 3 + + def test_grouped_frames_overlay_per_frame(self): + from matplotlib.animation import FuncAnimation + grouped = [[_line_frame(0.0), _line_frame(0.5)], + [_line_frame(1.0), _line_frame(1.5)]] + anim = anim_mod.animate(grouped, no_show=True) + assert isinstance(anim, FuncAnimation) + assert anim._save_count == 2 + + def test_multiblock_groups_equal_frame_indices(self): + from matplotlib.animation import FuncAnimation + grouped = [_line_frame(0.0), _line_frame(0.5)] + anim = anim_mod.animate(grouped, multiblock=True, no_show=True) + assert isinstance(anim, FuncAnimation) + assert anim._save_count == 1 + + def test_grouptags_builds_one_animation_per_tag(self): + from matplotlib.animation import FuncAnimation + frames = _three_frames() + frames[0].tag = "left" + frames[1].tag = "left" + frames[2].tag = "right" + animations = anim_mod.animate(frames, grouptags=True, no_show=True) + assert len(animations) == 2 + assert all(isinstance(anim, FuncAnimation) for anim in animations) + assert [anim._save_count for anim in animations] == [2, 1] + + def test_show_true_does_not_raise_on_agg(self): + anim = anim_mod.animate(_three_frames(), no_show=False) + assert anim is not None + + @needs_ffmpeg + @external_tool + @slow + def test_live_animation_saves_mp4(self, tmp_path): + out = tmp_path / "live.mp4" + anim = anim_mod.animate(_three_frames(), + save=True, + saveas=str(out), + fps=5, + no_show=True) + assert anim is not None + assert out.exists() + assert out.stat().st_size > 0 + + def test_notitle_suppresses_frame_time_title(self): + fig = plt.figure() + anim_mod._render_frame(0, anim_mod._normalize_frames(_three_frames()), fig, + {"notitle": True}) + assert fig._suptitle is None + + def test_title_includes_frame_and_time_by_default(self): + fig = plt.figure() + anim_mod._render_frame(1, anim_mod._normalize_frames(_three_frames()), fig, + {}) + assert "frame: 1" in fig._suptitle.get_text() + assert "time:" in fig._suptitle.get_text() + + def test_explicit_title_is_not_clobbered_by_the_auto_title(self): + fig = plt.figure() + anim_mod._render_frame(1, anim_mod._normalize_frames(_three_frames()), fig, + {"title": "My Animation"}) + assert fig._suptitle.get_text() == "My Animation" + + @pytest.mark.parametrize(("ctx", "expected"), [ + ({ + "time": 1.25 + }, "time: 1.2500e+00"), + ({ + "frame": 7 + }, "frame: 7"), + ]) + def test_generated_title_accepts_either_frame_metadata_field( + self, ctx, expected): + frame = _line_frame(0.0) + frame.ctx.pop("frame") + frame.ctx.pop("time") + frame.ctx.update(ctx) + fig = plt.figure() + anim_mod._draw_frame([frame], fig, {}) + assert fig._suptitle.get_text() == expected + + def test_variable_range_skips_global_limit_calculation(self, monkeypatch): + monkeypatch.setattr( + anim_mod, "_frame_value_range", + lambda *_args, **_kwargs: pytest.fail("global range was calculated")) + anim = anim_mod.animate(_three_frames(), variable_range=True, no_show=True) + assert anim is not None + + def test_live_save_configuration_without_running_a_writer( + self, monkeypatch, tmp_path): + saved = [] + + def save(self, filename, **kwargs): + saved.append((filename, kwargs)) + + monkeypatch.setattr(anim_mod, "require_ffmpeg", lambda _caller: "/ffmpeg") + monkeypatch.setattr("matplotlib.animation.FuncAnimation.save", save) + out = tmp_path / "movie.mp4" + anim_mod.animate(_three_frames(), + save=True, + saveas=str(out), + fps=5, + dpi=80, + no_show=True) + assert saved == [(str(out), {"writer": "ffmpeg", "fps": 5, "dpi": 80})] + + +# -------------------------------------------------------------------------- +# saved frames +# -------------------------------------------------------------------------- + + +class TestSaveFrames: + + def test_writes_one_png_per_frame(self, tmp_path): + prefix = str(tmp_path / "frame") + paths = anim_mod.animate(_three_frames(), saveframes=prefix, no_show=True) + assert len(paths) == 3 + for p in paths: + assert os.path.isfile(p) + + def test_saveframes_path_naming(self, tmp_path): + prefix = str(tmp_path / "myframe") + paths = anim_mod.animate(_three_frames(), saveframes=prefix, no_show=True) + assert paths[0] == f"{prefix}_0.png" + assert paths[2] == f"{prefix}_2.png" + + def test_nproc_parallel_writes_the_same_frames(self, tmp_path): + prefix = str(tmp_path / "frame") + paths = anim_mod.animate(_three_frames(), + saveframes=prefix, + nproc=2, + no_show=True) + assert len(paths) == 3 + for p in paths: + assert os.path.isfile(p) + + def test_nproc_without_saveframes_compiles_through_a_scratch_dir( + self, tmp_path): + out = tmp_path / "parallel.gif" + result = anim_mod.animate(_three_frames(), + nproc=2, + tmpdir=str(tmp_path), + saveas=str(out), + no_show=True) + assert result == str(out) + assert out.exists() + # the scratch directory must not leak its frame PNGs behind. + assert list(tmp_path.glob("*.png")) == [] + + def test_worker_can_render_one_frame_directly(self, tmp_path): + prefix = str(tmp_path / "worker") + path = anim_mod._save_frame_worker( + (3, [_line_frame(0.0)], {}, prefix, 72, (3.0, 2.0))) + assert path == f"{prefix}_3.png" + assert os.path.isfile(path) + + def test_grouped_tags_suffix_saved_frame_prefixes(self, tmp_path): + frames = _three_frames() + frames[0].tag = "left" + frames[1].tag = "left" + frames[2].tag = "right" + prefix = str(tmp_path / "frame") + paths = anim_mod.animate(frames, + grouptags=True, + saveframes=prefix, + no_show=True) + assert paths == [[f"{prefix}_left_0.png", f"{prefix}_left_1.png"], + [f"{prefix}_right_0.png"]] + + def test_saveas_without_extension_defaults_to_gif(self, tmp_path): + prefix = str(tmp_path / "frame") + output = tmp_path / "movie" + anim_mod.animate(_three_frames(), + saveframes=prefix, + saveas=str(output), + no_show=True) + assert output.with_suffix(".gif").is_file() + + +# -------------------------------------------------------------------------- +# movie compile +# -------------------------------------------------------------------------- + + +class TestCompileMovie: + + def test_unsupported_extension_raises(self, tmp_path): + with pytest.raises(ValueError, match="unsupported"): + anim_mod._compile_movie([], str(tmp_path / "out.bogus"), duration=100.0) + + def test_gif_compile_via_pil(self, tmp_path): + prefix = str(tmp_path / "frame") + paths = anim_mod.animate(_three_frames(), saveframes=prefix, no_show=True) + out = tmp_path / "out.gif" + anim_mod._compile_movie(paths, str(out), duration=100.0) + assert out.exists() + + def test_animate_saves_gif_end_to_end(self, tmp_path): + out = tmp_path / "movie.gif" + prefix = str(tmp_path / "frame") + result = anim_mod.animate(_three_frames(), + saveframes=prefix, + save=True, + saveas=str(out), + no_show=True) + assert out.exists() + assert len(result) == 3 + + def test_video_extension_raises_clearly_without_ffmpeg( + self, monkeypatch, tmp_path): + monkeypatch.setattr(_ffmpeg, "resolve_ffmpeg", lambda: None) + prefix = str(tmp_path / "frame") + paths = anim_mod.animate(_three_frames(), saveframes=prefix, no_show=True) + with pytest.raises(RuntimeError, match="ffmpeg"): + anim_mod._compile_movie(paths, str(tmp_path / "out.mp4"), duration=100.0) + + @pytest.mark.parametrize("fail_encoding", [False, True]) + def test_video_writer_protocol_without_external_process( + self, monkeypatch, tmp_path, fail_encoding): + from contextlib import contextmanager + from PIL import Image + import matplotlib.animation + + events = [] + images = [] + + class FakeImage: + size = (320, 200) + + def __init__(self): + self.closed = False + images.append(self) + + def __enter__(self): + return self + + def __exit__(self, *_exc): + self.closed = True + + class FakeAxes: + + def axis(self, value): + events.append(("axis", value)) + + def clear(self): + events.append(("clear", )) + + def imshow(self, image): + assert isinstance(image, FakeImage) + assert not image.closed + events.append(("imshow", )) + + class FakeFigure: + + def add_axes(self, bounds): + assert bounds == [0, 0, 1, 1] + return FakeAxes() + + class FakeWriter: + + def __init__(self, fps): + events.append(("fps", fps)) + + @contextmanager + def saving(self, figure, output_file, dpi): + assert isinstance(figure, FakeFigure) + events.append(("saving", output_file, dpi)) + yield + + def grab_frame(self): + events.append(("grab", )) + if fail_encoding: + raise RuntimeError("encoding failed") + + monkeypatch.setattr(anim_mod, "require_ffmpeg", lambda _caller: "/ffmpeg") + monkeypatch.setattr(Image, "open", lambda _path: FakeImage()) + monkeypatch.setattr(matplotlib.animation, "FFMpegWriter", FakeWriter) + monkeypatch.setattr(plt, "figure", lambda **_kwargs: FakeFigure()) + monkeypatch.setattr(plt, "close", lambda figure: events.append( + ("close", figure))) + + output = str(tmp_path / "movie.mp4") + if fail_encoding: + with pytest.raises(RuntimeError, match="encoding failed"): + anim_mod._compile_movie(["one.png", "two.png"], output, duration=250.0) + else: + anim_mod._compile_movie(["one.png", "two.png"], output, duration=250.0) + assert ("fps", 4.0) in events + assert ("saving", output, 100) in events + assert events.count(("grab", )) == (1 if fail_encoding else 2) + assert len(images) == (2 if fail_encoding else 3) + assert all(image.closed for image in images) + assert events[-1][0] == "close" + + @needs_ffmpeg + @external_tool + @slow + def test_mp4_compile_with_ffmpeg(self, tmp_path): + prefix = str(tmp_path / "frame") + paths = anim_mod.animate(_three_frames(), saveframes=prefix, no_show=True) + out = tmp_path / "out.mp4" + anim_mod._compile_movie(paths, str(out), fps=10, duration=100.0) + assert out.exists() + assert out.stat().st_size > 0 diff --git a/tests/test_render_ffmpeg.py b/tests/test_render_ffmpeg.py new file mode 100644 index 00000000..8b116119 --- /dev/null +++ b/tests/test_render_ffmpeg.py @@ -0,0 +1,48 @@ +"""Tests for postgkyl.render._ffmpeg -- shared ffmpeg discovery.""" + +from __future__ import annotations + +import builtins + +import pytest + +from postgkyl.render import _ffmpeg + + +def test_resolve_prefers_path(monkeypatch): + monkeypatch.setattr(_ffmpeg.shutil, "which", lambda _name: "/usr/bin/ffmpeg") + assert _ffmpeg.resolve_ffmpeg() == "/usr/bin/ffmpeg" + + +def test_resolve_falls_back_to_imageio_ffmpeg(monkeypatch): + import types + + fake_module = types.SimpleNamespace( + get_ffmpeg_exe=lambda: "/fake/imageio_ffmpeg/ffmpeg") + monkeypatch.setattr(_ffmpeg.shutil, "which", lambda _name: None) + monkeypatch.setitem(__import__("sys").modules, "imageio_ffmpeg", fake_module) + assert _ffmpeg.resolve_ffmpeg() == "/fake/imageio_ffmpeg/ffmpeg" + + +def test_resolve_returns_none_when_both_sources_are_missing(monkeypatch): + real_import = builtins.__import__ + + def missing_imageio(name, *args, **kwargs): + if name == "imageio_ffmpeg": + raise ImportError(name) + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(_ffmpeg.shutil, "which", lambda _name: None) + monkeypatch.setattr(builtins, "__import__", missing_imageio) + assert _ffmpeg.resolve_ffmpeg() is None + + +def test_require_raises_clearly_when_nothing_resolves(monkeypatch): + monkeypatch.setattr(_ffmpeg, "resolve_ffmpeg", lambda: None) + with pytest.raises(RuntimeError, match="ffmpeg"): + _ffmpeg.require_ffmpeg("animate") + + +def test_require_returns_resolved_path(monkeypatch): + monkeypatch.setattr(_ffmpeg, "resolve_ffmpeg", lambda: "/usr/bin/ffmpeg") + assert _ffmpeg.require_ffmpeg("animate") == "/usr/bin/ffmpeg" diff --git a/tests/test_render_labels.py b/tests/test_render_labels.py new file mode 100644 index 00000000..de34b32e --- /dev/null +++ b/tests/test_render_labels.py @@ -0,0 +1,60 @@ +"""Tests for postgkyl.render.labels -- latex_to_unicode / latex_to_html.""" + +from __future__ import annotations + +from postgkyl.render.labels import latex_to_html, latex_to_unicode + + +class TestLatexToUnicode: + + def test_empty_string_passthrough(self): + assert latex_to_unicode("") == "" + + def test_plain_text_unchanged(self): + assert latex_to_unicode("hello") == "hello" + + def test_strips_dollar_delimiters(self): + assert latex_to_unicode(r"$\mu$") == "μ" + + def test_greek_letter_without_dollars(self): + assert latex_to_unicode(r"\rho") == "ρ" + + def test_multiple_greek_letters(self): + assert latex_to_unicode(r"\alpha \beta \gamma") == "α β γ" + + def test_uppercase_greek_letters(self): + assert latex_to_unicode( + r"\Omega \Delta \Theta \Sigma \Lambda") == "Ω Δ Θ Σ Λ" + + def test_parallel_and_perp_with_subscripts_unconverted(self): + assert latex_to_unicode(r"$\mu_{\parallel}$") == "μ_{∥}" + assert latex_to_unicode(r"E_{\perp}") == "E_{⊥}" + + def test_strips_surrounding_whitespace(self): + assert latex_to_unicode(" \\pi ") == "π" + + +class TestLatexToHtml: + + def test_empty_string_passthrough(self): + assert latex_to_html("") == "" + + def test_plain_text_unchanged(self): + result = latex_to_html("field") + assert result == "field" + + def test_brace_subscript_becomes_html_sub(self): + result = latex_to_html(r"$B_{x}$") + assert result == "Bx" + + def test_bare_subscript_becomes_html_sub(self): + result = latex_to_html("n_0") + assert result == "n0" + + def test_greek_letter_converted(self): + result = latex_to_html(r"$\omega$") + assert "ω" in result + + def test_greek_and_subscript_combined(self): + assert latex_to_html(r"$\mu_{\parallel}$") == "μ∥" + assert latex_to_html(r"E_{\perp}") == "E⊥" diff --git a/tests/test_render_matplotlib.py b/tests/test_render_matplotlib.py new file mode 100644 index 00000000..ea324e06 --- /dev/null +++ b/tests/test_render_matplotlib.py @@ -0,0 +1,807 @@ +"""Tests for postgkyl.render.matplotlib -- multi-panel figures, the pgkyl +colorbar, log axes, vmin/vmax, aspect, and mapped (curvilinear) grids. + +``render.plot``'s basic single/multi-dataset 1-D and 2-D behaviour is already +covered by ``tests/test_coverage_leaf.py`` and ``tests/test_postgkyl.py``; +this file focuses on the features layer 09 adds on top. +""" + +from __future__ import annotations + +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython, operations +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.render import matplotlib as backend + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +GEN = os.path.join(DATA, "generated") + + +def _line(n=8, offset=0.0) -> GDataState: + d = GDataState() + d.push([np.linspace(0.0, 1.0, n + 1)], + (np.arange(n, dtype=float) + offset)[:, None]) + return d + + +def _field_2d(n=8, ncomp=1) -> GDataState: + d = GDataState() + grid = [np.linspace(0.0, 1.0, n + 1), np.linspace(0.0, 1.0, n + 1)] + values = np.stack([ + np.arange(n * n, dtype=float).reshape(n, n) + 10.0 * c + for c in range(ncomp) + ], + axis=-1) + d.push(grid, values) + return d + + +@pytest.fixture(autouse=True) +def _close_figs(): + plt.close("all") + yield + plt.close("all") + + +# -------------------------------------------------------------------------- +# Multi-panel (multi-component) layout +# -------------------------------------------------------------------------- + + +class TestMultiPanel: + + def test_two_components_get_two_axes(self): + fig = backend.plot(_field_2d(ncomp=2), no_show=True) + assert len(fig.axes) >= 2 + + def test_four_components_use_a_square_grid(self): + fig = backend.plot(_field_2d(ncomp=4), no_show=True) + # 4 components -> 2x2 grid -> 4 panel axes, each with its own colorbar axes. + assert len(fig.axes) == 8 + + def test_five_components_hides_the_leftover_axis(self): + fig = backend.plot(_field_2d(ncomp=5), no_show=True) + off_axes = [ax for ax in fig.axes if not ax.axison] + assert len(off_axes) == 1 + + def test_single_component_has_no_per_panel_title(self): + fig = backend.plot(_field_2d(ncomp=1), no_show=True) + assert fig.axes[0].get_title() == "" + + def test_legend_can_be_limited_to_one_subplot_and_relocated(self): + a = _line() + b = _line(offset=3.0) + a.values = np.column_stack((a.values[:, 0], a.values[:, 0] + 1.0)) + b.values = np.column_stack((b.values[:, 0], b.values[:, 0] + 1.0)) + + fig = backend.plot(a, + b, + multiblock=True, + no_show=True, + legend_labels=["first", "second"], + legend_subplot=1, + legend_loc="lower left") + + assert fig.axes[0].get_legend() is None + legend = fig.axes[1].get_legend() + assert legend is not None + assert legend._loc == 3 # Matplotlib's code for "lower left". + assert [text.get_text() + for text in legend.get_texts()] == ["first", "second"] + + def test_legend_subplot_rejects_an_out_of_range_index(self): + with pytest.raises(ValueError, match="between 0 and 0"): + backend.plot(_line(), no_show=True, legend_subplot=1) + + +# -------------------------------------------------------------------------- +# The pgkyl colorbar +# -------------------------------------------------------------------------- + + +class TestColorbar: + + def test_colorbar_true_adds_an_axes(self): + fig = backend.plot(_field_2d(), no_show=True, no_colorbar=False) + assert len(fig.axes) == 2 # the panel + the appended colorbar axes + + def test_colorbar_false_omits_it(self): + fig = backend.plot(_field_2d(), no_show=True, no_colorbar=True) + assert len(fig.axes) == 1 + + def test_clabel_reaches_the_colorbar(self): + fig = backend.plot(_field_2d(), + no_show=True, + no_colorbar=False, + clabel="density") + cbar_ax = fig.axes[1] + assert cbar_ax.get_ylabel() == "density" + + +# -------------------------------------------------------------------------- +# Log axes +# -------------------------------------------------------------------------- + + +class TestLogAxes: + + def test_logx_1d(self): + fig = backend.plot(_line(), no_show=True, logx=True) + assert fig.axes[0].get_xscale() == "log" + + def test_logy_1d(self): + fig = backend.plot(_line(), no_show=True, logy=True) + assert fig.axes[0].get_yscale() == "log" + + def test_logz_uses_lognorm_on_2d_colormap(self): + d = _field_2d() + d.values[...] = d.values + 1.0 # keep strictly positive for LogNorm + fig = backend.plot(d, no_show=True, logz=True) + im = fig.axes[0].collections[0] + from matplotlib.colors import LogNorm + assert isinstance(im.norm, LogNorm) + + +# -------------------------------------------------------------------------- +# Grid indices +# -------------------------------------------------------------------------- + + +class TestGridIndices: + + def test_1d_uses_zero_based_indices_instead_of_grid_values(self): + data = GDataState() + time = np.array([0.0, 0.1, 0.4, 1.2]) + data.push([time], np.arange(time.size, dtype=float)[:, None]) + + fig = backend.plot(data, no_show=True, grid_indices=True) + + np.testing.assert_array_equal(fig.axes[0].lines[0].get_xdata(), + np.arange(time.size)) + assert fig.get_supxlabel() == r"$i_0$" + np.testing.assert_array_equal(data.grid[0], time) + + def test_2d_puts_cell_centers_at_integer_indices(self): + data = _field_2d(n=4) + + fig = backend.plot(data, no_show=True, grid_indices=True, no_colorbar=True) + + coordinates = fig.axes[0].collections[0].get_coordinates() + assert coordinates[..., 0].min() == pytest.approx(-0.5) + assert coordinates[..., 0].max() == pytest.approx(3.5) + assert coordinates[..., 1].min() == pytest.approx(-0.5) + assert coordinates[..., 1].max() == pytest.approx(3.5) + assert fig.get_supxlabel() == r"$i_0$" + assert fig.get_supylabel() == r"$i_1$" + + +# -------------------------------------------------------------------------- +# Joined linear/log split panels +# -------------------------------------------------------------------------- + + +class TestSplitLinearLog: + + @staticmethod + def _split_line(ncomp=1): + d = GDataState() + # Cell centers are exactly [-2, -1, 0, 1, 2], pinning the split-point + # ownership rule (the point itself belongs to the right panel). + grid = [np.linspace(-2.5, 2.5, 6)] + base = np.array([1.0, 2.0, 3.0, 10.0, 100.0]) + values = np.stack([base * (comp + 1) for comp in range(ncomp)], axis=-1) + d.push(grid, values) + return d + + def test_each_component_becomes_a_linear_left_log_right_pair(self): + fig = backend.plot(self._split_line(ncomp=2), + no_show=True, + split_linear_log=True) + + assert len(fig.axes) == 4 + assert [axis.get_yscale() + for axis in fig.axes] == ["linear", "log", "linear", "log"] + np.testing.assert_allclose(fig.axes[0].lines[0].get_xdata(), [-2.0, -1.0]) + np.testing.assert_allclose(fig.axes[1].lines[0].get_xdata(), + [0.0, 1.0, 2.0]) + assert fig.axes[0].get_xlim()[1] == pytest.approx(0.0) + assert fig.axes[1].get_xlim()[0] == pytest.approx(0.0) + + def test_split_point_and_log_side_are_configurable(self): + fig = backend.plot(self._split_line(), + no_show=True, + split_linear_log=True, + split_point=1.0, + split_log_side="left", + split_log_base=2) + + left, right = fig.axes + assert left.get_yscale() == "log" + assert right.get_yscale() == "linear" + assert left.yaxis._scale.base == 2 + np.testing.assert_allclose(left.lines[0].get_xdata(), [-2.0, -1.0, 0.0]) + np.testing.assert_allclose(right.lines[0].get_xdata(), [1.0, 2.0]) + + def test_per_component_linear_and_log_limits(self): + fig = backend.plot(self._split_line(ncomp=2), + no_show=True, + split_linear_log=True, + split_linear_ylim=[(0.0, 5.0), (-2.0, 8.0)], + split_log_ylim={ + 0: (1.0, 200.0), + 1: (2.0, 400.0) + }) + + assert fig.axes[0].get_ylim() == (0.0, 5.0) + assert fig.axes[1].get_ylim() == (1.0, 200.0) + assert fig.axes[2].get_ylim() == (-2.0, 8.0) + assert fig.axes[3].get_ylim() == (2.0, 400.0) + + def test_shared_limit_pair_applies_to_every_component(self): + fig = backend.plot(self._split_line(ncomp=2), + no_show=True, + split_linear_log=True, + split_linear_ylim=(0.0, 10.0)) + assert fig.axes[0].get_ylim() == (0.0, 10.0) + assert fig.axes[2].get_ylim() == (0.0, 10.0) + + def test_legend_subplot_uses_log_half_of_logical_subplot(self): + a = self._split_line(ncomp=2) + b = self._split_line(ncomp=2) + fig = backend.plot(a, + b, + multiblock=True, + no_show=True, + split_linear_log=True, + legend_labels=["a", "b"], + legend_subplot=1, + split_legend_side="log") + + assert fig.axes[0].get_legend() is None + assert fig.axes[1].get_legend() is None + assert fig.axes[2].get_legend() is None + legend = fig.axes[3].get_legend() + assert [text.get_text() for text in legend.get_texts()] == ["a", "b"] + + def test_pair_geometry_and_logical_labels(self): + fig = backend.plot(self._split_line(), + no_show=True, + split_linear_log=True, + split_width_ratios=(2.0, 1.0), + split_gap=0.05, + subplot_titles="density", + subplot_ylabels="n", + subplot_xlabels="z") + left, right = fig.axes + assert left.get_position().width / right.get_position( + ).width == pytest.approx(2.0) + assert left.get_title() == "density" + assert left.get_ylabel() == "n" + assert left.get_xlabel() == "z" + assert right.yaxis.get_ticks_position() == "right" + + def test_left_owns_seam_tick_label_by_default(self): + fig = backend.plot(self._split_line(), no_show=True, split_linear_log=True) + left_ticks = fig.axes[0].get_xticks() + right_ticks = fig.axes[1].get_xticks() + assert left_ticks[-1] == pytest.approx(0.0) + assert right_ticks[0] > 0.0 + + @pytest.mark.parametrize("kwargs, message", [ + ({ + "split_log_side": "middle" + }, "split_log_side"), + ({ + "split_legend_side": "middle" + }, "split_legend_side"), + ({ + "split_width_ratios": (1.0, 0.0) + }, "split_width_ratios"), + ({ + "split_gap": -0.1 + }, "split_gap"), + ({ + "split_log_nonpositive": "drop" + }, "split_log_nonpositive"), + ({ + "split_seam_ticklabels": "middle" + }, "split_seam_ticklabels"), + ({ + "split_log_base": 1.0 + }, "split_log_base"), + ]) + def test_invalid_split_options_raise(self, kwargs, message): + with pytest.raises(ValueError, match=message): + backend.plot(self._split_line(), + no_show=True, + split_linear_log=True, + **kwargs) + + def test_split_rejects_2d_transpose_and_logy(self): + with pytest.raises(ValueError, match="only supported for 1D"): + backend.plot(_field_2d(), no_show=True, split_linear_log=True) + with pytest.raises(ValueError, match="transpose"): + backend.plot(self._split_line(), + no_show=True, + split_linear_log=True, + transpose=True) + with pytest.raises(ValueError, match="logy"): + backend.plot(self._split_line(), + no_show=True, + split_linear_log=True, + logy=True) + + +# -------------------------------------------------------------------------- +# value ranges: ymin/ymax (1-D), zmin/zmax (2-D color range) +# -------------------------------------------------------------------------- + + +class TestValueRange: + + def test_ymin_ymax_set_1d_ylim(self): + fig = backend.plot(_line(), no_show=True, ymin=-5.0, ymax=50.0) + assert fig.axes[0].get_ylim() == (-5.0, 50.0) + + def test_zmin_zmax_set_2d_colormap_range(self): + fig = backend.plot(_field_2d(), no_show=True, zmin=0.0, zmax=1.0) + im = fig.axes[0].collections[0] + assert im.get_clim() == (0.0, 1.0) + + +# -------------------------------------------------------------------------- +# Aspect +# -------------------------------------------------------------------------- + + +class TestAspect: + + def test_aspect_applies_to_2d_axes(self): + # aspect only takes effect with fixaspect=True -- --aspect on the CLI + # implies --fix-aspect (see cli/commands/plot.py), but the render engine + # itself keeps the two independent, exactly as main's output.plot did. + fig = backend.plot(_field_2d(), no_show=True, fixaspect=True, aspect=1.0) + assert fig.axes[0].get_aspect() == 1.0 + + def test_aspect_none_leaves_default(self): + fig = backend.plot(_field_2d(), no_show=True) + assert fig.axes[0].get_aspect() == "auto" + + +# -------------------------------------------------------------------------- +# cmap / diverging +# -------------------------------------------------------------------------- + + +class TestColormap: + + def test_explicit_cmap_is_used(self): + fig = backend.plot(_field_2d(), no_show=True, cmap="plasma") + im = fig.axes[0].collections[0] + assert im.get_cmap().name == "plasma" + + def test_diverging_uses_rdbu(self): + fig = backend.plot(_field_2d(), no_show=True, diverging=True) + im = fig.axes[0].collections[0] + assert im.get_cmap().name == "RdBu_r" + + +# -------------------------------------------------------------------------- +# style / rcParams +# -------------------------------------------------------------------------- + + +class TestStyleAndRcParams: + + def test_style_kwarg_applies_named_style(self): + backend.plot(_line(), no_show=True, style="default") + import matplotlib as mpl + assert mpl.rcParams["image.cmap"] == "viridis" + + def test_rcparams_dict_overrides(self): + backend.plot(_line(), no_show=True, rcParams={"lines.linewidth": 5.0}) + import matplotlib as mpl + assert mpl.rcParams["lines.linewidth"] == 5.0 + + +# -------------------------------------------------------------------------- +# output +# -------------------------------------------------------------------------- + + +class TestSaving: + + @pytest.mark.parametrize("extension", [".png", ".pdf"]) + def test_saveas_writes_supported_formats(self, tmp_path, extension): + output = tmp_path / f"figure{extension}" + fig = backend.plot(_line(), no_show=True, saveas=output) + + assert output.exists() + assert fig is plt.gcf() + + def test_extensionless_saveas_defaults_to_png(self, tmp_path): + output = tmp_path / "figure" + backend.plot(_line(), no_show=True, saveas=output) + + assert (tmp_path / "figure.png").exists() + + def test_empty_saveas_is_inert(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + backend.plot(_line(), no_show=True, saveas="") + assert list(tmp_path.iterdir()) == [] + + def test_save_true_derives_a_png_name(self, tmp_path, monkeypatch): + data = _line() + data._file_name = "/input/run.gkyl" + monkeypatch.chdir(tmp_path) + + backend.plot(data, no_show=True, save=True) + + assert (tmp_path / "run.png").exists() + + def test_saveas_sequence_writes_each_requested_format(self, tmp_path): + outputs = [tmp_path / "figure.png", tmp_path / "figure.pdf"] + backend.plot(_line(), no_show=True, saveas=outputs) + assert all(output.exists() for output in outputs) + + def test_unsupported_save_extension_raises(self, tmp_path): + with pytest.raises(ValueError, match="Supported formats are: .png, .pdf"): + backend.plot(_line(), no_show=True, saveas=tmp_path / "figure.svg") + + +# -------------------------------------------------------------------------- +# fig reuse (the hook render.animate needs) +# -------------------------------------------------------------------------- + + +class TestFigureReuse: + + def test_reusing_a_figure_clears_previous_axes(self): + fig = plt.figure() + backend.plot(_line(), no_show=True, figure=fig, clear=True) + first_axes_id = id(fig.axes[0]) + backend.plot(_line(offset=5.0), no_show=True, figure=fig, clear=True) + assert len(fig.axes) == 1 + assert id(fig.axes[0]) != first_axes_id + + +# -------------------------------------------------------------------------- +# transpose -- swap the horizontal and vertical axes (upstream PR #225) +# -------------------------------------------------------------------------- + + +def _field_2d_rect(n0=4, n1=8) -> GDataState: + d = GDataState() + grid = [np.linspace(0.0, 1.0, n0 + 1), np.linspace(0.0, 2.0, n1 + 1)] + values = np.arange(n0 * n1, dtype=float).reshape(n0, n1)[..., None] + d.push(grid, values) + return d + + +class TestTranspose: + + def test_1d_puts_the_coordinate_on_the_vertical_axis(self): + fig = backend.plot(_line(), no_show=True, transpose=True) + line = fig.axes[0].lines[0] + edges = np.linspace(0.0, 1.0, 9) + np.testing.assert_allclose(line.get_ydata(), 0.5 * (edges[:-1] + edges[1:])) + np.testing.assert_allclose(line.get_xdata(), np.arange(8, dtype=float)) + + def test_1d_default_label_follows_the_coordinate(self): + fig = backend.plot(_line(), no_show=True, transpose=True) + assert fig.get_supylabel() == r"$z_0$" + assert fig.get_supxlabel() == "" + + def test_2d_swaps_the_mesh_axes(self): + n0, n1 = 4, 8 + d = _field_2d_rect(n0, n1) + fig = backend.plot(d, no_show=True, transpose=True) + im = fig.axes[0].collections[0] + # The horizontal axis now carries dimension 1 (extent 0..2), the + # vertical dimension 0 (extent 0..1); the quads' value layout follows. + assert im.get_coordinates().shape == (n0 + 1, n1 + 1, 2) + np.testing.assert_allclose(fig.axes[0].get_xlim(), (0.0, 2.0)) + np.testing.assert_allclose(fig.axes[0].get_ylim(), (0.0, 1.0)) + np.testing.assert_allclose( + np.asarray(im.get_array()).reshape(n0, n1), d.values[..., 0]) + + def test_2d_swaps_the_default_labels(self): + fig = backend.plot(_field_2d_rect(), no_show=True, transpose=True) + assert fig.get_supxlabel() == r"$z_1$" + assert fig.get_supylabel() == r"$z_0$" + + def test_2d_does_not_mutate_the_dataset(self): + d = _field_2d_rect() + cells_before = d.num_cells.copy() + values_before = d.values.copy() + backend.plot(d, no_show=True, transpose=True) + np.testing.assert_array_equal(d.num_cells, cells_before) + np.testing.assert_array_equal(d.values, values_before) + + +# -------------------------------------------------------------------------- +# Mapped (curvilinear) grids -- MAPPING.md's BACKEND row +# -------------------------------------------------------------------------- + + +@needs_gkeyll +class TestMappedGrids: + + def test_2d_curvilinear_grid_plots_via_pcolormesh(self): + data = pg.load(os.path.join(GEN, "2d_ms_p1.gkyl")).interpolate() + mapped = operations.map(data, + os.path.join(GEN, "2d_c2p_stretch_ms_p1.gkyl"), + space="conf") + assert mapped.grid[0].ndim == 2 # genuinely curvilinear + fig = mapped.plot(no_show=True) + assert fig is not None + im = fig.axes[0].collections[0] + assert im.get_array().size > 0 + + def test_1d_non_uniform_mapped_axis_uses_true_centers(self): + """A 1-D vel map produces non-uniform edges; _centers must handle them + generically (it already does -- this pins the behaviour).""" + edges = np.array([0.0, 1.0, 4.0, 9.0, 16.0]) # non-uniform, monotone + d = GDataState() + d.push([edges], np.arange(4, dtype=float)[:, None]) + fig = backend.plot(d, no_show=True) + line = fig.axes[0].lines[0] + x_plotted = line.get_xdata() + np.testing.assert_allclose(x_plotted, 0.5 * (edges[:-1] + edges[1:])) + + +# -------------------------------------------------------------------------- +# surface plots +# -------------------------------------------------------------------------- + + +class TestSurface: + + def test_surface_uses_3d_axes(self): + fig = backend.plot(_field_2d(), no_show=True, surface=True) + assert fig.axes[0].name == "3d" + + def test_surface_without_comparison_gets_a_colorbar(self): + fig = backend.plot(_field_2d(), no_show=True, surface=True) + assert len(fig.axes) == 2 # the 3D panel + its colorbar + + def test_surface_alpha_is_applied(self): + fig = backend.plot(_field_2d(), no_show=True, surface=True, alpha=0.3) + poly3d = fig.axes[0].collections[0] + assert poly3d.get_alpha() == pytest.approx(0.3) + + +# -------------------------------------------------------------------------- +# multi-dataset 2D overlay comparison (surface/contour) +# -------------------------------------------------------------------------- + + +class TestComparisonOverlay: + + def test_contour_comparison_gives_each_dataset_its_own_color_and_legend(self): + fig = backend.plot(_field_2d(), + _field_2d(), + multiblock=True, + no_show=True, + contour=True, + comparison=True, + legend_labels=["a", "b"]) + ax = fig.axes[0] + assert ax.get_legend() is not None + handles = ax.get_legend().legend_handles + assert len(handles) == 2 + assert handles[0].get_facecolor() != handles[1].get_facecolor() + + def test_surface_comparison_gives_each_dataset_its_own_color_and_legend(self): + fig = backend.plot(_field_2d(), + _field_2d(), + multiblock=True, + no_show=True, + surface=True, + comparison=True, + legend_labels=["a", "b"]) + ax = fig.axes[0] + assert ax.get_legend() is not None + assert len(ax.get_legend().legend_handles) == 2 + + +# -------------------------------------------------------------------------- +# cval-based colormap coloring for 1D lines +# -------------------------------------------------------------------------- + + +class TestCvalColoring: + + def test_line_colored_by_cval(self): + fig = backend.plot(_line(), + no_show=True, + cmap="viridis", + cval=0.0, + cval_min=0.0, + cval_max=1.0) + assert fig.axes[0].lines[0].get_color() == plt.get_cmap("viridis")(0.0) + + def test_second_call_into_the_same_figure_uses_its_own_cval(self): + fig = backend.plot(_line(), + no_show=True, + cmap="viridis", + cval=0.0, + cval_min=0.0, + cval_max=1.0) + backend.plot(_line(offset=1), + figure=fig, + no_show=True, + cmap="viridis", + cval=1.0, + cval_min=0.0, + cval_max=1.0) + colors = [line.get_color() for line in fig.axes[0].lines] + assert colors[0] == plt.get_cmap("viridis")(0.0) + assert colors[1] == plt.get_cmap("viridis")(1.0) + + def test_cval_without_cmap_is_ignored(self): + fig = backend.plot(_line(), no_show=True, cval=0.5, color="red") + assert fig.axes[0].lines[0].get_color() == "red" + + +# -------------------------------------------------------------------------- +# Explicit colors for 1D lines +# -------------------------------------------------------------------------- + + +class TestLineColors: + + def test_color_sequence_assigns_one_color_to_each_dataset(self): + fig = backend.plot(_line(), + _line(offset=1), + _line(offset=2), + multiblock=True, + no_show=True, + color=["tab:red", "tab:green", "tab:blue"]) + + assert [line.get_color() for line in fig.axes[0].lines + ] == ["tab:red", "tab:green", "tab:blue"] + + def test_scalar_color_still_applies_to_every_line(self): + fig = backend.plot(_line(), + _line(offset=1), + multiblock=True, + no_show=True, + color="purple") + assert [line.get_color() + for line in fig.axes[0].lines] == ["purple", "purple"] + + def test_dataset_colors_repeat_across_component_panels(self): + a = _line() + b = _line(offset=2) + a.values = np.column_stack((a.values[:, 0], a.values[:, 0] + 1)) + b.values = np.column_stack((b.values[:, 0], b.values[:, 0] + 1)) + + fig = backend.plot(a, + b, + multiblock=True, + no_show=True, + color=["red", "blue"]) + + assert [line.get_color() for line in fig.axes[0].lines] == ["red", "blue"] + assert [line.get_color() for line in fig.axes[1].lines] == ["red", "blue"] + + def test_color_sequence_follows_dataset_then_component_order(self): + a = _line() + b = _line(offset=2) + a.values = np.column_stack((a.values[:, 0], a.values[:, 0] + 1)) + b.values = np.column_stack((b.values[:, 0], b.values[:, 0] + 1)) + + fig = backend.plot(a, + b, + multiblock=True, + no_show=True, + color=["red", "orange", "blue", "cyan"]) + + assert [line.get_color() for line in fig.axes[0].lines] == ["red", "blue"] + assert [line.get_color() + for line in fig.axes[1].lines] == ["orange", "cyan"] + + def test_rgb_tuple_remains_a_single_color(self): + rgb = (0.1, 0.2, 0.3) + fig = backend.plot(_line(), + _line(offset=1), + multiblock=True, + no_show=True, + color=rgb) + assert [line.get_color() for line in fig.axes[0].lines] == [rgb, rgb] + + def test_color_sequence_length_must_match_line_count(self): + with pytest.raises(ValueError, match="2 entries.*expected either 3.*or 3"): + backend.plot(_line(), + _line(offset=1), + _line(offset=2), + multiblock=True, + no_show=True, + color=["red", "blue"]) + + +# -------------------------------------------------------------------------- +# Per-dataset linestyles for 1D lines +# -------------------------------------------------------------------------- + + +class TestLineStyles: + + def test_linestyle_sequence_assigns_one_style_to_each_dataset(self): + fig = backend.plot(_line(), + _line(offset=1), + multiblock=True, + no_show=True, + linestyle=["-", "--"]) + + assert [line.get_linestyle() for line in fig.axes[0].lines] == ["-", "--"] + + def test_scalar_linestyle_applies_to_every_dataset(self): + fig = backend.plot(_line(), + _line(offset=1), + multiblock=True, + no_show=True, + linestyle=":") + assert [line.get_linestyle() for line in fig.axes[0].lines] == [":", ":"] + + def test_single_entry_sequence_applies_to_every_dataset(self): + fig = backend.plot(_line(), + _line(offset=1), + multiblock=True, + no_show=True, + linestyle=["-."]) + assert [line.get_linestyle() for line in fig.axes[0].lines] == ["-.", "-."] + + def test_omitted_linestyle_does_not_override_plot_format(self): + fig = backend.plot(_line(), no_show=True, args=["--"]) + assert fig.axes[0].lines[0].get_linestyle() == "--" + + def test_dataset_linestyles_repeat_across_component_panels(self): + a = _line() + b = _line(offset=2) + a.values = np.column_stack((a.values[:, 0], a.values[:, 0] + 1)) + b.values = np.column_stack((b.values[:, 0], b.values[:, 0] + 1)) + + fig = backend.plot(a, + b, + multiblock=True, + no_show=True, + linestyle=["-", "--"]) + + assert [line.get_linestyle() for line in fig.axes[0].lines] == ["-", "--"] + assert [line.get_linestyle() for line in fig.axes[1].lines] == ["-", "--"] + + def test_dataset_linestyles_apply_to_both_split_axes(self): + fig = backend.plot(_line(), + _line(offset=1), + multiblock=True, + no_show=True, + linestyle=["-", "--"], + split_linear_log=True, + split_point=0.5) + + for axis in fig.axes: + assert [line.get_linestyle() for line in axis.lines] == ["-", "--"] + + def test_linestyle_sequence_length_must_match_dataset_count(self): + with pytest.raises(ValueError, match="2 entries.*expected either 1.*or 3"): + backend.plot(_line(), + _line(offset=1), + _line(offset=2), + multiblock=True, + no_show=True, + linestyle=["-", "--"]) diff --git a/tests/test_render_matplotlib_coverage.py b/tests/test_render_matplotlib_coverage.py new file mode 100644 index 00000000..30aded59 --- /dev/null +++ b/tests/test_render_matplotlib_coverage.py @@ -0,0 +1,705 @@ +"""Coverage top-up for postgkyl.render.matplotlib. + +``tests/test_render_matplotlib.py`` covers the layer's headline features; +this file targets the branches ``--cov-report=term-missing`` still flagged: +``_nodal_grid``'s error/curvilinear paths, the rcParams novelties, label +shift/scale formatting, figure creation/reuse edge cases, contour/quiver/ +streamline/lineouts, zmin/zmax ``extend``, logz diverging, and the 0-D +data path. No compiled Gkeyll is needed -- every dataset here is built by +hand with ``GDataState.push``. +""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.render import matplotlib as backend + + +def _line(n=8, offset=0.0) -> GDataState: + d = GDataState() + d.push([np.linspace(0.0, 1.0, n + 1)], + (np.arange(n, dtype=float) + offset)[:, None]) + return d + + +def _field_2d(n=8, ncomp=1) -> GDataState: + d = GDataState() + grid = [np.linspace(0.0, 1.0, n + 1), np.linspace(0.0, 1.0, n + 1)] + values = np.stack([ + np.arange(n * n, dtype=float).reshape(n, n) + 10.0 * c + for c in range(ncomp) + ], + axis=-1) + d.push(grid, values) + return d + + +class _UnlabelledState: + + _file_name = "" + + +@pytest.fixture(autouse=True) +def _close_figs(): + plt.close("all") + yield + plt.close("all") + + +# -------------------------------------------------------------------------- +# Pure output/style normalization helpers +# -------------------------------------------------------------------------- + + +class TestOutputNormalization: + + def test_indexed_saveas_handles_one_path_and_a_sequence(self): + assert backend._indexed_saveas("plot.pdf", 3, True) == "plot_3.pdf" + assert backend._indexed_saveas(("a", "b.png"), 2, + True) == ("a_2", "b_2.png") + + def test_default_output_stem_uses_labels_and_fallbacks(self): + labelled = _line() + labelled.label = "ion density" + assert backend._default_output_stem([labelled, _UnlabelledState() + ]) == "ion_density_dataset_1" + assert backend._default_output_stem([]) == "matplotlib_output" + + def test_output_paths_rejects_a_non_iterable(self): + with pytest.raises(TypeError, match="path or an iterable"): + backend._output_paths(False, 3, []) + + def test_output_paths_rejects_a_non_path_entry(self): + with pytest.raises(TypeError, match="entry must be path-like"): + backend._output_paths(False, ["ok.png", object()], []) + + +class TestLineStyleNormalization: + + def test_invalid_color_string_is_left_to_matplotlib(self): + assert backend._normalize_line_colors("not-a-color") is None + + def test_non_iterable_color_is_treated_as_a_scalar(self): + assert backend._normalize_line_colors(object()) is None + + def test_empty_color_sequence_raises(self): + with pytest.raises(ValueError, match="must not be an empty sequence"): + backend._normalize_line_colors([]) + + def test_invalid_color_sequence_entry_raises(self): + with pytest.raises(ValueError, match="every entry"): + backend._normalize_line_colors(["red", "not-a-color"]) + + def test_custom_dash_pattern_is_one_linestyle(self): + assert backend._normalize_linestyles((0, (5, 2)), 2) is None + + def test_non_iterable_linestyle_is_treated_as_a_scalar(self): + assert backend._normalize_linestyles(object(), 2) is None + + def test_empty_linestyle_sequence_raises(self): + with pytest.raises(ValueError, match="must not be an empty sequence"): + backend._normalize_linestyles([], 2) + + +class TestSmallPlotHelpers: + + def test_xkcd_without_a_suitable_font_warns_and_uses_sans_serif( + self, monkeypatch): + monkeypatch.setattr(backend.fm.fontManager, "ttflist", []) + with pytest.warns(UserWarning, match="No xkcd-style font"): + context, rc = backend.get_xkcd_safely() + assert callable(context) + assert rc == {"font.family": "sans-serif"} + + def test_shared_range_skips_missing_components_and_nonfinite_values(self): + one_comp = _field_2d(n=2) + two_comps = _field_2d(n=2, ncomp=2) + two_comps._values[..., 1] = np.nan + assert backend._shared_component_range([one_comp, two_comps], 0.0, + 1.0) == [(0.0, 3.0), (None, None)] + + @pytest.mark.parametrize(("limits", "comp", "expected"), + [({}, 0, None), ([(0.0, 1.0)], 2, None), + ([None], 0, None)]) + def test_split_ylim_missing_component_is_automatic(self, limits, comp, + expected): + assert backend._split_ylim_for_component(limits, comp) is expected + + def test_split_ylim_rejects_a_non_container(self): + with pytest.raises(TypeError, match="split y-limits"): + backend._split_ylim_for_component(1.0, 0) + + @pytest.mark.parametrize("limits", [[(0.0, 1.0, 2.0)], {0: 1.0}]) + def test_split_ylim_rejects_a_malformed_pair(self, limits): + with pytest.raises(ValueError, match="must be a .* pair"): + backend._split_ylim_for_component(limits, 0) + + +# -------------------------------------------------------------------------- +# _nodal_grid, tested directly as the pure function it is +# -------------------------------------------------------------------------- + + +class TestNodalGridDirect: + + def test_dim_count_mismatch_raises(self): + with pytest.raises(ValueError, match="doesn't match"): + backend._nodal_grid([np.linspace(0.0, 1.0, 5)], np.array([4, 4])) + + def test_1d_bad_edge_count_raises(self): + with pytest.raises(ValueError, match="terribly wrong"): + backend._nodal_grid([np.linspace(0.0, 1.0, 5)], np.array([10])) + + def test_1d_curvilinear_edges_averaged(self): + g = np.array([[0.0, 1.0], [2.0, 4.0]]) # shape (2, 2): cells[0]+1 == 2 + out = backend._nodal_grid([g], np.array([1])) + np.testing.assert_allclose(out[0], 0.5 * (g[:-1] + g[1:])) + + def test_2d_curvilinear_cell_centered_passthrough(self): + g0 = np.ones((3, 3)) + out = backend._nodal_grid([g0, np.ones((3, 3))], np.array([3, 3])) + assert out[0] is g0 + + def test_2d_curvilinear_edges_averaged(self): + g0 = np.arange(16, dtype=float).reshape(4, 4) # cells[0]+1 == 4 + out = backend._nodal_grid([g0, np.ones((4, 4))], np.array([3, 3])) + np.testing.assert_allclose(out[0], 0.5 * (g0[:-1, :-1] + g0[1:, 1:])) + + def test_2d_curvilinear_bad_shape_raises(self): + g0 = np.ones((5, 5)) + with pytest.raises(ValueError, match="terribly wrong"): + backend._nodal_grid([g0, np.ones((5, 5))], np.array([3, 3])) + + +# -------------------------------------------------------------------------- +# rcParams novelties: jet / xkcd / color / linewidth / linestyle +# -------------------------------------------------------------------------- + + +class TestRcParamNovelties: + + def test_jet_sets_cmap(self): + with mpl.rc_context(): + backend.plot(_field_2d(), no_show=True, jet=True) + assert mpl.rcParams["image.cmap"] == "jet" + + @pytest.mark.filterwarnings("ignore:No xkcd-style font found:UserWarning") + def test_xkcd_flag_invokes_xkcd_mode(self): + with mpl.rc_context(): + fig = backend.plot(_line(), no_show=True, xkcd=True) + line = fig.axes[0].lines[0] + assert line.get_sketch_params() is not None + + @pytest.mark.filterwarnings("ignore:No xkcd-style font found:UserWarning") + def test_xkcd_flag_does_not_leak_into_global_rcparams(self): + # A past bug: `plt.xkcd()` called without a `with` block never reverted, + # contaminating every plot drawn afterwards. + with mpl.rc_context(): + backend.plot(_line(), no_show=True, xkcd=True) + assert mpl.rcParams["path.sketch"] is None + + def test_color_sets_rcparam(self): + with mpl.rc_context(): + backend.plot(_line(), no_show=True, color="red") + assert mpl.rcParams["lines.color"] == "red" + + def test_linewidth_sets_rcparam(self): + with mpl.rc_context(): + backend.plot(_line(), no_show=True, linewidth=4.0) + assert mpl.rcParams["lines.linewidth"] == 4.0 + + def test_linestyle_sets_rcparam(self): + with mpl.rc_context(): + backend.plot(_line(), no_show=True, linestyle="--") + assert mpl.rcParams["lines.linestyle"] == "--" + + +# -------------------------------------------------------------------------- +# xlabel/ylabel/clabel shift-scale annotation branches +# -------------------------------------------------------------------------- + + +class TestLabelShiftScale: + + def test_xlabel_shift_and_scale(self): + fig = backend.plot(_line(), + no_show=True, + squeeze=True, + xshift=1.0, + xscale=2.0) + lbl = fig.axes[0].get_xlabel() + assert " + " in lbl and r"\times" in lbl + + def test_xlabel_shift_only(self): + fig = backend.plot(_line(), no_show=True, squeeze=True, xshift=1.0) + lbl = fig.axes[0].get_xlabel() + assert " + " in lbl and r"\times" not in lbl + + def test_xlabel_scale_only(self): + fig = backend.plot(_line(), no_show=True, squeeze=True, xscale=2.0) + lbl = fig.axes[0].get_xlabel() + assert r"\times" in lbl and " + " not in lbl + + def test_ylabel_shift_and_scale(self): + fig = backend.plot(_field_2d(), + no_show=True, + squeeze=True, + yshift=1.0, + yscale=2.0) + lbl = fig.axes[0].get_ylabel() + assert " + " in lbl and r"\times" in lbl + + def test_ylabel_bug_branch_uses_xshift(self): + fig = backend.plot(_field_2d(), no_show=True, squeeze=True, xshift=1.0) + lbl = fig.axes[0].get_ylabel() + assert " + " in lbl + + def test_ylabel_bug_branch_uses_xscale(self): + fig = backend.plot(_field_2d(), no_show=True, squeeze=True, xscale=2.0) + lbl = fig.axes[0].get_ylabel() + assert r"\times" in lbl + + def test_clabel_gets_zscale_annotation(self): + fig = backend.plot(_field_2d(), + no_show=True, + clabel="density", + zscale=2.0, + no_colorbar=False) + cbar_lbl = fig.axes[1].get_ylabel() + assert "density" in cbar_lbl and r"\times" in cbar_lbl + + +# -------------------------------------------------------------------------- +# figsize / figure kwarg / figure reuse +# -------------------------------------------------------------------------- + + +class TestFigureCreation: + + def test_figsize_string_is_parsed(self): + fig = backend.plot(_line(), no_show=True, figsize="6,4") + np.testing.assert_allclose(fig.get_size_inches(), (6.0, 4.0)) + + def test_figure_int_selects_numbered_figure(self): + fig = backend.plot(_line(), no_show=True, figure=11) + assert fig.number == 11 + + def test_figure_str_selects_numbered_figure(self): + fig = backend.plot(_line(), no_show=True, figure="12") + assert fig.number == 12 + + def test_figure_object_is_used_directly(self): + fig_obj = plt.figure() + result = backend.plot(_line(), no_show=True, figure=fig_obj) + assert result is fig_obj + assert len(result.axes) == 1 + + def test_figure_invalid_type_raises(self): + with pytest.raises(TypeError, match="'figure' keyword"): + backend.plot(_line(), no_show=True, figure=3.14) + + def test_reused_figure_without_enough_axes_raises(self): + fig_obj = plt.figure() + fig_obj.subplots(1, 1) + with pytest.raises(ValueError, match="not enough axes"): + backend.plot(_field_2d(ncomp=4), no_show=True, figure=fig_obj) + + +# -------------------------------------------------------------------------- +# squeeze=True single-panel layout +# -------------------------------------------------------------------------- + + +class TestSqueezeLayout: + + def test_squeeze_sets_title_on_first_use(self): + fig = backend.plot(_field_2d(ncomp=1), + no_show=True, + squeeze=True, + title="hello", + no_colorbar=True) + assert fig.axes[0].get_title() == "hello" + assert len(fig.axes) == 1 + + +# -------------------------------------------------------------------------- +# Multi-panel subplot titles +# -------------------------------------------------------------------------- + + +class TestSubplotTitles: + + def test_per_panel_titles_are_set(self): + fig = backend.plot(_field_2d(ncomp=2), + no_show=True, + no_colorbar=True, + subplot_titles="a,b") + assert fig.axes[0].get_title() == "a" + assert fig.axes[1].get_title() == "b" + + +# -------------------------------------------------------------------------- +# Multi-dataset label_prefix via data.get_label() +# -------------------------------------------------------------------------- + + +class TestMultiDatasetLabel: + + def test_label_prefix_uses_dataset_label(self): + a, b = _line(), _line(offset=3.0) + a.label, b.label = "first", "second" + fig = backend.plot(a, b, multiblock=True, no_show=True) + texts = [t.get_text() for t in fig.axes[0].get_legend().get_texts()] + assert "first" in texts and "second" in texts + + +# -------------------------------------------------------------------------- +# Dimensionality errors raised per-dataset (not just from the first/"ref") +# -------------------------------------------------------------------------- + + +class TestDimensionalityErrors: + + def test_second_dataset_over_2d_raises(self): + ok = _line() + bad = GDataState() + bad.push([np.linspace(0, 1, 3)] * 3, np.zeros((2, 2, 2, 1))) + with pytest.raises(ValueError, match="Only 1D and 2D"): + backend.plot(ok, bad, no_show=True) + + def test_0d_dataset_raises(self): + d = GDataState() + d.push([np.array([0.0, 1.0]), np.array([0.0, 1.0])], np.zeros((1, 1, 1))) + with pytest.raises(ValueError, match="0D data not supported"): + backend.plot(d, no_show=True) + + +# -------------------------------------------------------------------------- +# squeeze-with-curvilinear-grid dimension drop (plot()'s own inline squeeze) +# -------------------------------------------------------------------------- + + +class TestSqueezeCurvilinearDrop: + + def test_curvilinear_coordinate_meaned_over_dropped_axis(self): + d = GDataState() + g0 = np.array([[0.0], [1.0]]) + g1 = np.arange(24, dtype=float).reshape(4, + 6) # joint-shaped, no size-1 axis + values = np.arange(5, dtype=float).reshape(1, 5, 1) + d.push([g0, g1], values) + fig = backend.plot(d, no_show=True) + expected_edges = g1.mean(axis=0) + expected_x = 0.5 * (expected_edges[:-1] + expected_edges[1:]) + np.testing.assert_allclose(fig.axes[0].lines[0].get_xdata(), expected_x) + + +# -------------------------------------------------------------------------- +# contour +# -------------------------------------------------------------------------- + + +class TestContour: + + def test_default_levels_with_clabel_text(self): + fig = backend.plot(_field_2d(), no_show=True, contour=True, cont_label=True) + assert fig is not None + + def test_cnlevels_sets_integer_level_count(self): + fig = backend.plot(_field_2d(), no_show=True, contour=True, cnlevels=6) + assert fig is not None + + def test_clevels_colon_syntax_is_linspace(self): + fig = backend.plot(_field_2d(), + no_show=True, + contour=True, + clevels="0:60:5") + assert fig is not None + + def test_clevels_single_value_disables_colorbar(self): + fig = backend.plot(_field_2d(), no_show=True, contour=True, clevels="30") + assert len(fig.axes) == 1 + + def test_clevels_comma_list(self): + fig = backend.plot(_field_2d(), + no_show=True, + contour=True, + clevels="10,30,50") + assert fig is not None + + +# -------------------------------------------------------------------------- +# quiver / streamline (need a wide-enough grid so `skip` isn't 0) +# -------------------------------------------------------------------------- + + +class TestQuiverAndStreamline: + + def test_quiver_draws_vector_field(self): + fig = backend.plot(_field_2d(n=15, ncomp=2), no_show=True, quiver=True) + assert len(fig.axes[0].collections) >= 1 + + def test_quiver_on_curvilinear_grid_uses_2d_nodal_grid(self): + edges = np.linspace(0.0, 1.0, 16) + gx, gy = np.meshgrid(edges, edges, indexing="ij") + values = np.stack([np.zeros((15, 15)), np.ones((15, 15))], axis=-1) + d = GDataState() + d.push([gx, gy], values) + fig = backend.plot(d, no_show=True, quiver=True) + assert len(fig.axes[0].collections) >= 1 + + def test_streamline_default_uses_speed_as_color(self): + fig = backend.plot(_field_2d(n=15, ncomp=2), no_show=True, streamline=True) + assert fig is not None + + def test_streamline_explicit_color(self): + fig = backend.plot(_field_2d(n=15, ncomp=2), + no_show=True, + streamline=True, + color="black") + assert fig is not None + + +# -------------------------------------------------------------------------- +# lineouts +# -------------------------------------------------------------------------- + + +class TestLineouts: + + def test_lineouts_0_draws_one_line_per_column(self): + d = _field_2d(n=4) + fig = backend.plot(d, no_show=True, lineouts=0) + assert len(fig.axes[0].lines) == 4 + assert len(fig.axes) == 2 # panel + the appended lineout colorbar + + def test_lineouts_1_draws_one_line_per_row(self): + d = _field_2d(n=4) + fig = backend.plot(d, no_show=True, lineouts=1) + assert len(fig.axes[0].lines) == 4 + assert len(fig.axes) == 2 + + +# -------------------------------------------------------------------------- +# zmin/zmax -> colorbar `extend` +# -------------------------------------------------------------------------- + + +class TestExtend: + + def test_zmax_only_extends_max(self): + fig = backend.plot(_field_2d(), no_show=True, zmax=5.0) + im = fig.axes[0].collections[0] + assert im.colorbar.extend == "max" + + def test_zmin_only_extends_min(self): + fig = backend.plot(_field_2d(), no_show=True, zmin=5.0) + im = fig.axes[0].collections[0] + assert im.colorbar.extend == "min" + + +# -------------------------------------------------------------------------- +# plain pcolormesh: nodal-grid fallback when grid already matches cell count +# -------------------------------------------------------------------------- + + +class TestNodalGridFallback: + + def test_cell_centered_grid_falls_back_through_nodal_grid(self): + d = GDataState() + grid = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 1.0, 5)] + values = np.arange(25, dtype=float).reshape(5, 5, 1) + d.push(grid, values) + fig = backend.plot(d, no_show=True) + assert len(fig.axes[0].collections) == 1 + + +# -------------------------------------------------------------------------- +# logz + diverging -> SymLogNorm +# -------------------------------------------------------------------------- + + +class TestLogzDiverging: + + def test_logz_diverging_uses_symlognorm(self): + fig = backend.plot(_field_2d(), no_show=True, logz=True, diverging=True) + im = fig.axes[0].collections[0] + from matplotlib.colors import SymLogNorm + assert isinstance(im.norm, SymLogNorm) + + +# -------------------------------------------------------------------------- +# hashtag watermark +# -------------------------------------------------------------------------- + + +class TestHashtag: + + def test_hashtag_adds_text(self): + fig = backend.plot(_line(), no_show=True, hashtag=True) + texts = [t.get_text() for t in fig.axes[0].texts] + assert "#pgkyl" in texts + + +# -------------------------------------------------------------------------- +# xmin/xmax -> set_xlim +# -------------------------------------------------------------------------- + + +class TestXlim: + + def test_xmin_xmax_set_xlim(self): + fig = backend.plot(_line(), no_show=True, xmin=0.2, xmax=0.8) + assert fig.axes[0].get_xlim() == (0.2, 0.8) + + +# -------------------------------------------------------------------------- +# num_axes spreads multiple single-component datasets across separate panels +# -------------------------------------------------------------------------- + + +class TestNumAxesAcrossDatasets: + + def test_cur_start_axes_advances_between_datasets(self): + a, b = _field_2d(ncomp=1), _field_2d(ncomp=1) + fig = backend.plot(a, b, multiblock=True, no_show=True, num_axes=2) + assert len(fig.axes[0].collections) == 1 + assert len(fig.axes[1].collections) == 1 + + +# -------------------------------------------------------------------------- +# Remaining validation and specialized drawing branches +# -------------------------------------------------------------------------- + + +class TestRemainingValidationBranches: + + def test_color_sequence_rejects_a_2d_plot(self): + with pytest.raises(ValueError, match="only supported for 1D"): + backend.plot(_field_2d(), no_show=True, color=["red", "blue"]) + + @pytest.mark.parametrize(("kwargs", "error", "message"), + [({ + "split_point": object() + }, TypeError, "split_point"), + ({ + "split_point": np.inf + }, ValueError, "split_point"), + ({ + "split_log_base": object() + }, TypeError, "split_log_base"), + ({ + "split_gap": object() + }, TypeError, "split_gap")]) + def test_split_numeric_options_reject_invalid_values(self, kwargs, error, + message): + with pytest.raises(error, match=message): + backend.plot(_line(), no_show=True, split_linear_log=True, **kwargs) + + def test_legend_subplot_rejects_a_non_integer(self): + with pytest.raises(TypeError, match="must be an integer"): + backend.plot(_line(), no_show=True, legend_subplot="0") + + def test_second_dataset_over_2d_raises_inside_one_family(self): + bad = GDataState() + bad.push([np.linspace(0, 1, 3)] * 3, np.zeros((2, 2, 2, 1))) + with pytest.raises(ValueError, match="Only 1D and 2D"): + backend.plot(_line(), bad, multiblock=True, no_show=True) + + def test_second_dataset_must_match_split_dimensionality(self): + with pytest.raises(ValueError, match="every dataset must be 1D"): + backend.plot(_line(), + _field_2d(), + multiblock=True, + no_show=True, + split_linear_log=True) + + +class TestRemainingSplitBranches: + + @pytest.mark.parametrize(("side", "legend_axis"), [("left", 0), ("right", 1), + ("linear", 0)]) + def test_explicit_split_legend_side(self, side, legend_axis): + data = _line() + data.label = "curve" + fig = backend.plot(data, + no_show=True, + forcelegend=True, + split_linear_log=True, + split_legend_side=side) + assert fig.axes[legend_axis].get_legend() is not None + + def test_split_layout_accepts_y_label_only_and_hides_right_ticks(self): + fig = backend.plot(_line(), + no_show=True, + split_linear_log=True, + xlabel="", + ylabel="amplitude", + no_split_right_ticks=True) + assert fig.get_supxlabel() == "" + assert fig.get_supylabel() == "amplitude" + assert fig.axes[1].yaxis.get_ticks_position() != "right" + + def test_split_plot_allows_logarithmic_x_axis(self): + fig = backend.plot(_line(), + no_show=True, + split_linear_log=True, + split_point=0.5, + logx=True) + assert all(axis.get_xscale() == "log" for axis in fig.axes) + + +class TestRemainingSurfaceBranches: + + @staticmethod + def _mapped_field() -> GDataState: + coordinates = np.linspace(0.0, 1.0, 4) + gx, gy = np.meshgrid(coordinates, coordinates, indexing="ij") + data = GDataState() + data.push([gx, gy], np.arange(16, dtype=float).reshape(4, 4, 1)) + return data + + def test_transpose_transposes_joint_coordinate_arrays(self): + fig = backend.plot(self._mapped_field(), + no_show=True, + transpose=True, + no_colorbar=True) + assert len(fig.axes[0].collections) == 1 + + def test_surface_transposes_joint_coordinates_without_a_colorbar(self): + fig = backend.plot(self._mapped_field(), + no_show=True, + surface=True, + no_colorbar=True) + assert fig.axes[0].name == "3d" + assert len(fig.axes) == 1 + + def test_surface_applies_color_label_and_z_limits(self): + fig = backend.plot(_field_2d(), + no_show=True, + surface=True, + clabel="density", + zmin=1.0, + zmax=9.0) + assert fig.axes[0].get_zlabel() == "density" + assert fig.axes[0].get_zlim() == (1.0, 9.0) + + def test_unlabelled_surface_comparison_needs_no_legend_handle(self): + fig = backend.plot(_field_2d(), no_show=True, surface=True, comparison=True) + assert fig.axes[0].get_legend() is None + + def test_unlabelled_contour_comparison_needs_no_legend_handle(self): + fig = backend.plot(_field_2d(), no_show=True, contour=True, comparison=True) + assert fig.axes[0].get_legend() is None + + def test_cval_without_bounds_uses_colormap_midpoint(self): + fig = backend.plot(_line(), no_show=True, cmap="viridis", cval=3.0) + assert fig.axes[0].lines[0].get_color() == plt.get_cmap("viridis")(0.5) diff --git a/tests/test_render_plotly.py b/tests/test_render_plotly.py new file mode 100644 index 00000000..ba4d6a53 --- /dev/null +++ b/tests/test_render_plotly.py @@ -0,0 +1,775 @@ +"""Tests for postgkyl.render.plotly -- 2-D surfaces, 3-D volumes/scatter, +animation, and rotating-figure export. + +Adapted from ``tests_bak/test_plot.py``'s ``plotly`` cases: the old tests fed +``(grid, values)`` tuples straight into ``pg.output.plotly``; this layer's +``plotly()`` takes a :class:`~postgkyl.gdatastate.gdatastate.GDataState` instead (no +dual "GData or tuple" signature -- see PYTHON_PRINCIPLES.md #9), so every +case below builds one via ``GDataState().push(...)``. +""" + +from __future__ import annotations + +from importlib import import_module +import sys +from types import SimpleNamespace + +import matplotlib + +matplotlib.use("Agg") +import matplotlib as mpl +import numpy as np +import plotly.graph_objects as go +import pytest + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.render import _ffmpeg +from postgkyl.render.plotly import ( + plotly, + plotly_animate, + save_rotating_plotly_figure, +) +from postgkyl.render.plotly import ( + _log_colorbar_ticks, + _opacity_mapping, + _prepare_2d_coordinates, + _prepare_3d_coordinates, +) + +needs_ffmpeg = pytest.mark.skipif( + _ffmpeg.resolve_ffmpeg() is None, + reason="ffmpeg not found on PATH or via imageio-ffmpeg") +external_tool = pytest.mark.external_tool +slow = pytest.mark.slow + + +def _chrome_available() -> bool: + # Kaleido v1+ needs a real Chrome/Chromium binary (its own download via + # `kaleido_get_chrome` or a system install) -- without one, + # start_sync_server()'s background thread dies and to_image() hangs + # forever waiting on a server that never came up, rather than raising. + try: + from choreographer.browsers.chromium import Chromium + return Chromium.find_browser(skip_local=False) is not None + except Exception: + return False + + +needs_chrome = pytest.mark.skipif(not _chrome_available(), + reason="no Chrome/Chromium found for kaleido") + +# kaleido's Chrome subprocess (managed by the `choreographer` library under +# start_sync_server()/stop_sync_server()) has produced an intermittent, +# non-reproducible-on-Linux segfault during CPython's own interpreter +# finalization -- well after the whole pytest session has already passed -- +# only ever seen on macOS CI. Skip there rather than let it take down the +# whole pytest process; see skip_macos_animate_save in test_cli_commands.py +# for the same pattern applied to an analogous matplotlib/macOS crash. +skip_macos_chrome = pytest.mark.skipif( + sys.platform == "darwin", + reason="intermittent segfault during kaleido's Chrome subprocess " + "teardown at interpreter shutdown on macOS -- not reproducible " + "on Linux") + + +def _state(grid, values) -> GDataState: + d = GDataState() + d.push(list(grid), values) + return d + + +def _volume_3d(fn=lambda x, y, z: x + y + z, n=4): + grid = [ + np.linspace(0.0, 1.0, n), + np.linspace(0.0, 1.0, n), + np.linspace(0.0, 1.0, n) + ] + x, y, z = np.meshgrid(*grid, indexing="ij") + values = fn(x, y, z)[..., np.newaxis] + return _state(grid, values) + + +def _surface_2d(n=4, m=5): + grid = [np.linspace(0.0, 1.0, n), np.linspace(0.0, 1.0, m)] + x, y = np.meshgrid(*grid, indexing="ij") + values = (x + 2.0 * y)[..., np.newaxis] + return _state(grid, values) + + +class TestPlotlySurface2D: + + def test_returns_a_surface_trace(self): + fig = plotly(_surface_2d()) + assert isinstance(fig, go.Figure) + assert isinstance(fig.data[0], go.Surface) + + def test_surface_z_matches_values(self): + n, m = 4, 5 + grid = [np.linspace(0.0, 1.0, n), np.linspace(0.0, 1.0, m)] + x, y = np.meshgrid(*grid, indexing="ij") + fig = plotly(_surface_2d(n, m)) + np.testing.assert_allclose(fig.data[0].z, x + 2.0 * y) + + def test_axis_ranges_match_data_extent(self): + fig = plotly(_surface_2d()) + np.testing.assert_allclose(fig.layout.scene.xaxis.range, (0.0, 1.0)) + np.testing.assert_allclose(fig.layout.scene.yaxis.range, (0.0, 1.0)) + np.testing.assert_allclose(fig.layout.scene.zaxis.range, (0.0, 3.0)) + + def test_scatter_mode_rejected_for_surface(self): + with pytest.raises(ValueError, match="scatter"): + plotly(_surface_2d(), scatter=True) + + def test_surface_logc_applies_log_colorscale(self): + fig = plotly(_surface_2d(), logc=True, cmin=1.0e-3, cmax=10.0) + np.testing.assert_allclose(fig.data[0].cmin, -3.0) + np.testing.assert_allclose(fig.data[0].cmax, 1.0) + + def test_scale_and_shift_apply_to_surface_coordinates_and_height(self): + # x/y scale+shift the coordinates; z/color inherit from the *value* + # (zscale/zshift), matching src_bak/postgkyl/output/plotly.py:720. + n, m = 4, 5 + fig = plotly(_surface_2d(n, m), + xscale=2.0, + xshift=1.0, + yscale=3.0, + yshift=0.5, + zscale=2.0, + zshift=1.0) + np.testing.assert_allclose(fig.data[0].x.min(), 2.0) + np.testing.assert_allclose(fig.data[0].x.max(), 4.0) + np.testing.assert_allclose(fig.data[0].y.min(), 1.5) + np.testing.assert_allclose(fig.data[0].y.max(), 4.5) + np.testing.assert_allclose(np.nanmin(fig.data[0].z), 1.0) + np.testing.assert_allclose(np.nanmax(fig.data[0].z), 7.0) + + +class TestPlotly3DVolume: + + def test_returns_a_volume_trace_with_default_surface_count(self): + fig = plotly(_volume_3d()) + assert isinstance(fig, go.Figure) + assert fig.data[0].surface.count == 32 + + def test_axis_ranges_match_data_extent(self): + fig = plotly(_volume_3d()) + np.testing.assert_allclose(fig.layout.scene.xaxis.range, (0.0, 1.0)) + np.testing.assert_allclose(fig.layout.scene.yaxis.range, (0.0, 1.0)) + np.testing.assert_allclose(fig.layout.scene.zaxis.range, (0.0, 1.0)) + + def test_explicit_ranges_and_surface_count_override(self): + fig = plotly(_volume_3d(), + xrange=(0.2, 0.8), + yrange=(0.1, 0.9), + zrange=(0.3, 0.7), + surface_count=12) + np.testing.assert_allclose(fig.layout.scene.xaxis.range, (0.2, 0.8)) + np.testing.assert_allclose(fig.layout.scene.yaxis.range, (0.1, 0.9)) + np.testing.assert_allclose(fig.layout.scene.zaxis.range, (0.3, 0.7)) + assert fig.data[0].surface.count == 12 + + def test_color_scale_shift_and_clim(self): + fig = plotly(_volume_3d(), cscale=2.0, cshift=1.0, clim=(1.5, 5.5)) + np.testing.assert_allclose(fig.data[0].cmin, 1.5) + np.testing.assert_allclose(fig.data[0].cmax, 5.5) + np.testing.assert_allclose(np.nanmin(fig.data[0].value), 1.0) + np.testing.assert_allclose(np.nanmax(fig.data[0].value), 7.0) + + def test_logc_converts_linear_clim_to_log_space(self): + fig = plotly(_volume_3d(fn=lambda x, y, z: 1.0e-2 + x + y + z), + logc=True, + cmin=1.0e-20, + cmax=1.0e-2) + np.testing.assert_allclose(fig.data[0].cmin, -20.0) + np.testing.assert_allclose(fig.data[0].cmax, -2.0) + + def test_aspect_cube_mode(self): + fig = plotly(_volume_3d(), aspect="cube") + assert fig.layout.scene.aspectmode == "cube" + + def test_aspect_string_sets_mode(self): + fig = plotly(_volume_3d(), aspect="data") + assert fig.layout.scene.aspectmode == "data" + + def test_aspect_numeric_sets_manual_ratio(self): + fig = plotly(_volume_3d(), aspect=2.0) + assert fig.layout.scene.aspectmode == "manual" + assert fig.layout.scene.aspectratio.x == 2.0 + assert fig.layout.scene.aspectratio.y == 2.0 + assert fig.layout.scene.aspectratio.z == 2.0 + + def test_aspect_numeric_string_sets_manual_ratio(self): + fig = plotly(_volume_3d(), aspect="1.5") + assert fig.layout.scene.aspectmode == "manual" + assert fig.layout.scene.aspectratio.x == 1.5 + + def test_scale_and_shift_apply_to_volume_coordinates(self): + fig = plotly(_volume_3d(), + xscale=2.0, + xshift=1.0, + yscale=3.0, + yshift=0.5, + zscale=4.0, + zshift=1.0) + np.testing.assert_allclose(fig.layout.scene.xaxis.range, (2.0, 4.0)) + np.testing.assert_allclose(fig.layout.scene.yaxis.range, (1.5, 4.5)) + np.testing.assert_allclose(fig.layout.scene.zaxis.range, (4.0, 8.0)) + + def test_zscale_zshift_apply_to_volume_color_value(self): + # value = (x+y+z)*zscale + zshift, independent of the z *coordinate*'s + # own scale/shift -- matches src_bak/postgkyl/output/plotly.py:720. + fig = plotly(_volume_3d(), zscale=2.0, zshift=1.0) + np.testing.assert_allclose(np.nanmin(fig.data[0].value), 1.0) + np.testing.assert_allclose(np.nanmax(fig.data[0].value), 7.0) + + def test_cylindrical_to_cartesian_conversion(self): + r = np.linspace(0.0, 1.0, 4) + z = np.linspace(-0.5, 0.5, 4) + phi = np.linspace(0.0, 2.0 * np.pi, 5) + rr, zz, pp = np.meshgrid(r, z, phi, indexing="ij") + values = (rr + zz)[..., np.newaxis] + fig = plotly(_state([r, z, phi], values), cylindrical_to_cartesian=True) + np.testing.assert_allclose(fig.layout.scene.xaxis.range, (-1.0, 1.0), + atol=1e-12) + np.testing.assert_allclose(fig.layout.scene.yaxis.range, (-1.0, 1.0), + atol=1e-12) + np.testing.assert_allclose(fig.layout.scene.zaxis.range, (-0.5, 0.5), + atol=1e-12) + + +class TestPlotly3DScatter: + + def test_scatter_trace_basic_properties(self): + fig = plotly(_volume_3d(), + scatter=True, + marker_radius=3.0, + markerstyle="square", + cmin=0.2, + cmax=2.8) + assert isinstance(fig.data[0], go.Scatter3d) + assert fig.data[0].mode == "markers" + np.testing.assert_allclose(fig.data[0].marker.size, 6.0) + assert fig.data[0].marker.symbol == "square" + np.testing.assert_allclose(fig.data[0].marker.cmin, 0.2) + np.testing.assert_allclose(fig.data[0].marker.cmax, 2.8) + + def test_scatter_downsampling(self): + fig = plotly(_volume_3d(), scatter=True, maximum_points_per_axis=2) + # size-4 axis downsampled to indices [0, 2, 3] -> 3 points per axis. + assert len(fig.data[0].x) == 27 + assert len(fig.data[0].y) == 27 + assert len(fig.data[0].z) == 27 + + def test_opacity_gradient_when_requested(self): + fig = plotly(_volume_3d(), + scatter=True, + opacity=0.5, + scatter_opacity_range=(0.01, 1.0)) + colorscale = fig.data[0].marker.colorscale + low_alpha = float(colorscale[0][1].split(",")[-1].rstrip(")")) + high_alpha = float(colorscale[-1][1].split(",")[-1].rstrip(")")) + assert low_alpha < high_alpha + + def test_uniform_opacity_by_default(self): + fig = plotly(_volume_3d(), scatter=True, opacity=0.5) + colorscale = fig.data[0].marker.colorscale + low_alpha = float(colorscale[0][1].split(",")[-1].rstrip(")")) + high_alpha = float(colorscale[-1][1].split(",")[-1].rstrip(")")) + np.testing.assert_allclose(low_alpha, high_alpha) + np.testing.assert_allclose(fig.data[0].marker.opacity, 0.5) + + def test_log_opacity_ramp(self): + fig = plotly(_volume_3d(), + scatter=True, + scatter_opacity_range=(0.01, 1.0), + scatter_opacity_log=True) + colorscale = fig.data[0].marker.colorscale + alphas = np.array( + [float(c.split(",")[-1].rstrip(")")) for _, c in colorscale]) + q1 = int(0.25 * (len(alphas) - 1)) + q3 = int(0.75 * (len(alphas) - 1)) + low_span = alphas[q1] - alphas[0] + high_span = alphas[-1] - alphas[q3] + assert low_span > high_span + + +class TestPlotlyMultiComponent: + + def test_two_components_get_two_scenes(self): + grid = [ + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 4) + ] + x, y, z = np.meshgrid(*grid, indexing="ij") + values = np.stack([x + y + z, x - y - z], axis=-1) + fig = plotly(_state(grid, values)) + assert len(fig.data) == 2 + + def test_squeeze_forces_a_single_scene(self): + grid = [ + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 4) + ] + x, y, z = np.meshgrid(*grid, indexing="ij") + values = np.stack([x + y + z, x - y - z], axis=-1) + fig = plotly(_state(grid, values), squeeze=True) + assert len(fig.data) == 1 + + +class TestPlotlyMisc: + + def test_diverging_symmetric_colorscale(self): + fig = plotly(_volume_3d(), diverging=True) + assert fig.data[0].cmin == -fig.data[0].cmax + + def test_title_is_set(self): + fig = plotly(_volume_3d(), title="my title") + assert fig.layout.title.text == "my title" + + def test_hashtag_annotation(self): + fig = plotly(_volume_3d(), hashtag=True) + assert len(fig.layout.annotations) == 1 + assert fig.layout.annotations[0].text == "#pgkyl" + + def test_figsize_sets_pixel_dimensions(self): + fig = plotly(_volume_3d(), figsize=(6, 4)) + assert fig.layout.width == 600 + assert fig.layout.height == 400 + + def test_invalid_num_dims_raises(self): + d = _state([np.linspace(0.0, 1.0, 5)], np.ones((4, 1))) + with pytest.raises(ValueError, match="2D surface"): + plotly(d) + + def test_solid_color_disables_colorbar(self): + fig = plotly(_volume_3d(), color="red") + assert fig.data[0].showscale is False + + +class TestPlotlyStyleAndTheme: + + def test_light_background_sets_light_theme_colors(self): + fig = plotly(_volume_3d(), background="light") + assert fig.layout.paper_bgcolor == "#ffffff" + + def test_dark_background_is_the_default(self): + fig = plotly(_volume_3d()) + assert fig.layout.paper_bgcolor == "#000000" + + def test_explicit_style_kwarg_is_applied(self): + # "default" resets Matplotlib's baseline rc, distinct from the packaged + # postgkyl style's lines.linewidth == 2 (image.cmap gets overwritten + # right after by the cmap-resolution step below, so assert on a rc key + # that step never touches). + plotly(_volume_3d(), style="default") + assert mpl.rcParams["lines.linewidth"] == 1.5 + + def test_rcparams_override_is_applied(self): + plotly(_volume_3d(), rcParams={"lines.linewidth": 4.0}) + assert mpl.rcParams["lines.linewidth"] == 4.0 + + def test_invert_cmap_appends_reversal_suffix(self): + plotly(_volume_3d(), cmap="viridis", invert_cmap=True) + assert mpl.rcParams["image.cmap"] == "viridis_r" + + def test_invert_cmap_strips_reversal_suffix(self): + plotly(_volume_3d(), cmap="viridis_r", invert_cmap=True) + assert mpl.rcParams["image.cmap"] == "viridis" + + def test_xkcd_style_does_not_raise(self): + import matplotlib.pyplot as plt + plotly(_volume_3d(), xkcd=True) + plt.rcdefaults() + + +class TestPlotlyLogAxes: + + def test_log_axes_use_log10_ranges(self): + grid = [np.linspace(1.0, 10.0, 4), np.linspace(1.0, 100.0, 5)] + x, y = np.meshgrid(*grid, indexing="ij") + values = (x + y)[..., np.newaxis] + fig = plotly(_state(grid, values), logx=True, logy=True) + assert fig.layout.scene.xaxis.type == "log" + assert fig.layout.scene.yaxis.type == "log" + np.testing.assert_allclose(fig.layout.scene.xaxis.range, + [np.log10(1.0), np.log10(10.0)]) + + def test_logz_masks_nonpositive_volume_values(self): + # The z *coordinate* axis spans [0, 1] here, so log10(0) triggers an + # (expected, harmless) divide-by-zero warning independent of the + # *value* function -- match the old tree's behaviour, don't silence it + # at the source, just don't let it fail this test. + with np.errstate(divide="ignore"): + fig = plotly(_volume_3d(fn=lambda x, y, z: x + y + z - 1.4), logz=True) + # Values <= 0 become NaN in log space; the trace should still build. + assert isinstance(fig.data[0], go.Volume) + + def test_logc_with_all_nonpositive_values_uses_fallback_range(self): + fig = plotly(_volume_3d(fn=lambda x, y, z: -(x + y + z) - 1.0), logc=True) + assert isinstance(fig.data[0], go.Volume) + + def test_logc_ticks_append_max_when_step_overshoots_it(self): + # lo=0, hi=20 with the default max_ticks=7 steps by 3 and lands on 18, + # short of hi -- _log_colorbar_ticks must append the true endpoint. + fig = plotly(_surface_2d(), logc=True, cmin=1.0, cmax=1.0e20) + tick_vals = fig.data[0].colorbar.tickvals + assert tick_vals[-1] == 20.0 + + def test_logc_cmax_below_cmin_falls_back_to_a_one_decade_span(self): + # cmax < cmin collapses the requested log range; _apply_log_colorscale + # falls back to a single decade above cmin rather than an inverted one. + fig = plotly(_surface_2d(), logc=True, cmin=100.0, cmax=10.0) + np.testing.assert_allclose(fig.data[0].cmin, 2.0) + np.testing.assert_allclose(fig.data[0].cmax, 3.0) + + def test_all_nan_values_yield_nan_color_range_without_raising(self): + n, m = 4, 5 + grid = [np.linspace(0.0, 1.0, n), np.linspace(0.0, 1.0, m)] + values = np.full((n - 1, m - 1, 1), np.nan) + fig = plotly(_state(grid, values)) + assert np.isnan(fig.data[0].cmin) + assert np.isnan(fig.data[0].cmax) + + +class TestPlotlyPrivateHelpers: + """Direct tests for small pure helpers whose edge branches are defensive + code unreachable through ``plotly()``'s public contract: the coordinate + helpers only ever see grids matching the checked ``num_dims`` and 1-D + nodal axes (guaranteed by ``GDataState``), and ``_log_colorbar_ticks`` + only ever gets called with the already-finite range ``_apply_log_colorscale`` + computes. Testing these directly is simpler and more honest than + contriving a ``GDataState`` that violates those invariants.""" + + def test_opacity_mapping_swaps_inverted_bounds(self): + colorscale = [[0.0, "rgba(10, 20, 30, 1.000)"], + [1.0, "rgba(10, 20, 30, 1.000)"]] + out = _opacity_mapping(colorscale, min_alpha=0.9, max_alpha=0.1) + first_alpha = float(out[0][1].split(",")[-1].rstrip(")")) + last_alpha = float(out[-1][1].split(",")[-1].rstrip(")")) + np.testing.assert_allclose(first_alpha, 0.1) + np.testing.assert_allclose(last_alpha, 0.9) + + def test_opacity_mapping_passes_through_non_rgba_and_malformed_colors(self): + colorscale = [[0.0, "rgba(1, 2, 3)"], [1.0, "#ff0000"]] + out = _opacity_mapping(colorscale, min_alpha=0.0, max_alpha=1.0) + assert out == [[0.0, "rgba(1, 2, 3)"], [1.0, "#ff0000"]] + + def test_log_colorbar_ticks_empty_for_non_finite_bounds(self): + assert _log_colorbar_ticks(float("nan"), 5.0) == ([], []) + + def test_prepare_3d_coordinates_rejects_wrong_count(self): + with pytest.raises(ValueError, match="three coordinate arrays"): + _prepare_3d_coordinates((np.array([0.0]), np.array([0.0])), (1, )) + + def test_prepare_3d_coordinates_passes_through_already_meshed_arrays(self): + mesh = np.zeros((2, 2, 2)) + out = _prepare_3d_coordinates((mesh, mesh, mesh), mesh.shape) + assert out[0] is mesh and out[1] is mesh and out[2] is mesh + + def test_prepare_2d_coordinates_rejects_wrong_count(self): + with pytest.raises(ValueError, match="two coordinate arrays"): + _prepare_2d_coordinates((np.array([0.0]), ), (1, )) + + def test_prepare_2d_coordinates_passes_through_already_meshed_arrays(self): + mesh = np.zeros((2, 2)) + out = _prepare_2d_coordinates((mesh, mesh), mesh.shape) + assert out[0] is mesh and out[1] is mesh + + +class TestSaveRotatingPlotlyFigure: + + def _scene_fig(self): + # plotly() always calls fig.update_layout(scene=...), guaranteeing a + # real "scene" key in the layout (a bare go.Figure(go.Surface(...)) + # only gets one once actually rendered by a Plotly frontend). + return plotly(_volume_3d()) + + def test_bad_extension_raises(self): + with pytest.raises(ValueError, match=r"\.gif, \.mp4, or \.html"): + save_rotating_plotly_figure(self._scene_fig(), "out.bogus", 0.0, 10, 60.0, + 2.0) + + def test_nonpositive_fps_raises(self): + with pytest.raises(ValueError, match="fps"): + save_rotating_plotly_figure(self._scene_fig(), "out.gif", 0.0, 0, 60.0, + 2.0) + + def test_nonpositive_rotation_period_raises(self): + with pytest.raises(ValueError, match="rotation_period"): + save_rotating_plotly_figure(self._scene_fig(), "out.gif", 0.0, 10, 60.0, + 0.0) + + def test_requires_a_3d_scene_figure(self): + flat_fig = go.Figure(go.Scatter(x=[0, 1], y=[0, 1])) + with pytest.raises(ValueError, match="3D scene"): + save_rotating_plotly_figure(flat_fig, "out.gif", 0.0, 10, 60.0, 2.0) + + def test_html_export_embeds_rotation_script(self, tmp_path): + out = tmp_path / "out.html" + save_rotating_plotly_figure(self._scene_fig(), str(out), 45.0, 10, 60.0, + 2.0) + assert out.exists() + text = out.read_text(encoding="utf-8") + assert "PGKYL" in text or len(text) > 0 + + def test_html_export_zero_rotation_period_omits_script(self, tmp_path): + # rotation_period must stay positive (checked above), but omega is + # driven to exactly 0.0 via math.inf -- any finite (however huge) period + # still yields omega > 0.0 in float64 and takes the *other* branch. Pass + # every angle/period by keyword: the previous version of this test + # passed a huge value positionally where it actually landed in + # ``polar_angle`` (not ``rotation_period``, which stayed a normal 2.0), + # so it never drove omega to zero at all -- see C6. + import math + + out = tmp_path / "out.html" + save_rotating_plotly_figure(self._scene_fig(), + str(out), + starting_azimuthal_angle=0.0, + fps=10, + polar_angle=60.0, + rotation_period=math.inf, + radius=2.0) + assert out.exists() + assert "recomputeRotationParams" not in out.read_text(encoding="utf-8") + + @pytest.mark.parametrize(("extension", "command_marker"), [ + ("mp4", "yuv420p"), + ("gif", "palettegen"), + ]) + def test_binary_export_protocol_without_external_process( + self, monkeypatch, tmp_path, extension, command_marker): + import subprocess + plotly_module = import_module("postgkyl.render.plotly") + events = [] + + class FakeLayout: + + @staticmethod + def to_plotly_json(): + return {"scene": {}} + + class FakeFigure: + layout = FakeLayout() + + def update_layout(self, **kwargs): + events.append(("layout", kwargs)) + + def to_image(self, *, format): + assert format == "png" + return b"png" + + fake_kaleido = SimpleNamespace( + start_sync_server=lambda **kwargs: events.append(("start", kwargs)), + stop_sync_server=lambda **kwargs: events.append(("stop", kwargs)), + ) + commands = [] + monkeypatch.setitem(sys.modules, "kaleido", fake_kaleido) + monkeypatch.setattr(plotly_module, "require_ffmpeg", + lambda _caller: "/ffmpeg") + monkeypatch.setattr( + subprocess, "run", lambda command, **kwargs: commands.append( + (command, kwargs))) + + output = tmp_path / f"rotation.{extension}" + save_rotating_plotly_figure(FakeFigure(), + str(output), + starting_azimuthal_angle=15.0, + fps=2, + polar_angle=60.0, + rotation_period=1.0) + assert [event[0] for event in events].count("layout") == 2 + assert events[0][0] == "start" and events[-1][0] == "stop" + assert len(commands) == 1 + assert command_marker in " ".join(commands[0][0]) + assert commands[0][1]["check"] is True + + @needs_ffmpeg + @needs_chrome + @skip_macos_chrome + @external_tool + @slow + def test_gif_export_end_to_end(self, tmp_path): + # fps * rotation_period = 2 -- the minimum frame count that still + # exercises the multi-frame rotation loop (fewer, and the `max(2, ...)` + # floor in save_rotating_plotly_figure would hide fps/rotation_period + # from the frame count entirely). Each frame drives a real Kaleido + # render, so keeping this small matters for test runtime. + out = tmp_path / "out.gif" + save_rotating_plotly_figure(self._scene_fig(), str(out), 0.0, 2, 1.0, 1.0) + assert out.exists() + assert out.stat().st_size > 0 + + @needs_ffmpeg + @needs_chrome + @skip_macos_chrome + @external_tool + @slow + def test_mp4_export_end_to_end(self, tmp_path): + out = tmp_path / "out.mp4" + save_rotating_plotly_figure(self._scene_fig(), str(out), 0.0, 2, 1.0, 1.0) + assert out.exists() + assert out.stat().st_size > 0 + + +class TestPlotlyAnimate: + + def test_builds_frames_and_controls(self): + n = 4 + grid = [np.linspace(0.0, 1.0, n), np.linspace(0.0, 1.0, n)] + x, y = np.meshgrid(*grid, indexing="ij") + values0 = (x + 2.0 * y)[..., np.newaxis] + values1 = (x + 2.0 * y + 0.5)[..., np.newaxis] + fig = plotly_animate( + [_state(grid, values0), _state(grid, values1)], frame_duration=40) + assert isinstance(fig, go.Figure) + assert isinstance(fig.data[0], go.Surface) + assert len(fig.frames) == 1 + assert fig.frames[0].name == "1" + assert fig.layout.updatemenus[0].buttons[0].label == "Play" + + def test_requires_at_least_one_dataset(self): + with pytest.raises(ValueError, match="at least one"): + plotly_animate([]) + + def test_frame_labels_length_mismatch_raises(self): + with pytest.raises(ValueError, match="frame_labels"): + plotly_animate([_surface_2d(), _surface_2d()], frame_labels=["only one"]) + + def test_mismatched_trace_count_between_frames_raises(self): + grid = [ + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 4), + np.linspace(0.0, 1.0, 4) + ] + x, y, z = np.meshgrid(*grid, indexing="ij") + one_comp = _state(grid, (x + y + z)[..., np.newaxis]) + two_comp = _state(grid, np.stack([x + y + z, x - y - z], axis=-1)) + with pytest.raises(ValueError, match="same number of traces"): + plotly_animate([one_comp, two_comp]) + + def test_save_adds_html_extension_and_show_writes_preview( + self, monkeypatch, tmp_path): + plotly_module = import_module("postgkyl.render.plotly") + written = [] + opened = [] + monkeypatch.setattr(go.Figure, "write_html", + lambda _fig, path: written.append(str(path))) + monkeypatch.setattr(plotly_module, "open_preview", + lambda path: opened.append(str(path))) + + output = tmp_path / "animation" + plotly_animate([_surface_2d()], saveas=str(output)) + assert written[-1] == f"{output}.html" + + plotly_animate([_surface_2d()], show=True) + assert written[-1].endswith("plotly-animate_preview.html") + assert opened == [written[-1]] + + explicit_html = tmp_path / "animation.html" + plotly_animate([_surface_2d()], saveas=str(explicit_html), show=True) + assert written[-1] == str(explicit_html) + assert opened[-1] == str(explicit_html) + + +class TestOutputHelpers: + + def test_default_stem_prefers_file_then_label_then_fallback(self): + plotly_module = import_module("postgkyl.render.plotly") + data = _surface_2d() + data._file_name = "/tmp/simulation.gkyl" + assert plotly_module._default_output_stem(data) == "simulation" + data._file_name = "" + data._custom_label = "density" + assert plotly_module._default_output_stem(data) == "density" + data._custom_label = "" + assert plotly_module._default_output_stem(data) == "plotly_output" + + def test_write_output_dispatches_rotating_and_plain_formats( + self, monkeypatch, tmp_path): + plotly_module = import_module("postgkyl.render.plotly") + calls = [] + + class FakeFigure: + + def write_html(self, path): + calls.append(("html", path)) + + monkeypatch.setattr( + plotly_module, "save_rotating_plotly_figure", + lambda fig, path, **kwargs: calls.append(("rotate", path, kwargs))) + figure = FakeFigure() + rotating = plotly_module._write_plotly_output(figure, + str(tmp_path / "figure.html"), + starting_azimuthal_angle=0.0, + polar_angle=60.0, + rotation_period=2.0, + fps=10) + plain = plotly_module._write_plotly_output(figure, + str(tmp_path / "figure.png"), + starting_azimuthal_angle=0.0, + polar_angle=60.0, + rotation_period=2.0, + fps=10) + empty = plotly_module._write_plotly_output(figure, + "", + starting_azimuthal_angle=0.0, + polar_angle=60.0, + rotation_period=2.0, + fps=10) + assert rotating.endswith("figure.html") + assert plain.endswith("figure.html") + assert empty == ".html" + assert [call[0] for call in calls] == ["rotate", "html", "html"] + + def test_preview_sanitizes_names_and_open_preview_uses_file_uri( + self, monkeypatch, tmp_path): + plotly_module = import_module("postgkyl.render.plotly") + saved = [] + opened = [] + monkeypatch.setattr(plotly_module.tempfile, "gettempdir", + lambda: str(tmp_path)) + monkeypatch.setattr(plotly_module, "save_rotating_plotly_figure", + lambda _fig, path, **_kwargs: saved.append(path)) + monkeypatch.setattr(plotly_module.webbrowser, "open", + lambda uri: opened.append(uri)) + path = plotly_module._preview_plotly_figure(object(), + " !!! ", + starting_azimuthal_angle=0.0, + polar_angle=60.0, + rotation_period=2.0, + fps=10) + assert path.endswith("plotly_preview_preview.html") + assert saved == [path] + named_path = plotly_module._preview_plotly_figure( + object(), + "named", + starting_azimuthal_angle=0.0, + polar_angle=60.0, + rotation_period=2.0, + fps=10) + assert named_path.endswith("named_preview.html") + plotly_module.open_preview(path) + assert opened[0].startswith("file://") + + def test_plotly_save_and_show_dispatch_to_output_helpers(self, monkeypatch): + plotly_module = import_module("postgkyl.render.plotly") + calls = [] + monkeypatch.setattr( + plotly_module, "_write_plotly_output", + lambda _fig, path, **_kwargs: calls.append( + ("write", path)) or "saved.html") + monkeypatch.setattr( + plotly_module, "_preview_plotly_figure", + lambda _fig, stem, **_kwargs: calls.append( + ("preview", stem)) or "preview.html") + monkeypatch.setattr(plotly_module, "open_preview", + lambda path: calls.append(("open", path))) + + plotly(_surface_2d(), save=True) + plotly(_surface_2d(), show=True) + plotly(_surface_2d(), saveas="figure.html", show=True) + assert [call[0] + for call in calls] == ["write", "preview", "open", "write", "open"] diff --git a/tests/test_render_prep.py b/tests/test_render_prep.py new file mode 100644 index 00000000..20c2bcf9 --- /dev/null +++ b/tests/test_render_prep.py @@ -0,0 +1,237 @@ +"""Tests for postgkyl.render._prep -- the dataset -> plottable-array prep +shared by every render backend (formerly axis_and_grid_prep + load_plot_data).""" + +from __future__ import annotations + +import numpy as np + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.render._prep import ( + default_axis_labels, + format_axis_label, + prep_plot_data, + resolve_axis_labels, + squeeze_collapsed_axes, + subplot_grid, +) + +# -------------------------------------------------------------------------- +# default_axis_labels / format_axis_label +# -------------------------------------------------------------------------- + + +class TestDefaultAxisLabels: + + def test_returns_one_label_per_dim(self): + labels = default_axis_labels(3) + assert labels == [r"$z_0$", r"$z_1$", r"$z_2$"] + + def test_zero_dims_is_empty(self): + assert default_axis_labels(0) == [] + + +class TestFormatAxisLabel: + + def test_no_shift_no_scale_passthrough(self): + assert format_axis_label("x", 0.0, 1.0) == "x" + + def test_shift_only(self): + result = format_axis_label("x", 1.0, 1.0) + assert result == r"x + 1.00e+00" + + def test_scale_only(self): + result = format_axis_label("x", 0.0, 2.0) + assert result == r"x $\times$ 2.00e+00" + + def test_shift_and_scale(self): + result = format_axis_label("x", 1.0, 2.0) + assert result == r"(x + 1.00e+00) $\times$ 2.00e+00" + + +# -------------------------------------------------------------------------- +# resolve_axis_labels +# -------------------------------------------------------------------------- + + +class TestResolveAxisLabels: + + def test_defaults_for_2d(self): + xl, yl, zl, cl = resolve_axis_labels(xlabel=None, + ylabel=None, + zlabel=None, + clabel="", + num_dims=2) + assert xl == r"$z_0$" + assert yl == r"$z_1$" + + def test_1d_has_no_default_ylabel(self): + xl, yl, zl, cl = resolve_axis_labels(xlabel=None, + ylabel=None, + zlabel=None, + clabel="", + num_dims=1) + assert xl == r"$z_0$" + assert yl == "" + + def test_custom_labels_pass_through(self): + xl, yl, zl, cl = resolve_axis_labels(xlabel="myX", + ylabel="myY", + zlabel="myZ", + clabel="myC", + num_dims=2, + zscale=2.0) + assert xl == "myX" + assert yl == "myY" + assert "2.00" in cl + + def test_3d_zlabel_defaults_to_third_axis(self): + xl, yl, zl, cl = resolve_axis_labels(xlabel=None, + ylabel=None, + zlabel=None, + clabel="", + num_dims=3) + assert zl == r"$z_2$" + + def test_clabel_annotated_with_zscale(self): + _, _, _, cl = resolve_axis_labels(xlabel=None, + ylabel=None, + zlabel=None, + clabel="density", + num_dims=2, + zscale=3.0) + assert cl == r"density $\times$ 3.000e+00" + + def test_clabel_zscale_with_no_base_label(self): + _, _, _, cl = resolve_axis_labels(xlabel=None, + ylabel=None, + zlabel=None, + clabel="", + num_dims=2, + zscale=3.0) + assert cl == r"$\times$ 3.000e+00" + + +# -------------------------------------------------------------------------- +# squeeze_collapsed_axes +# -------------------------------------------------------------------------- + + +class TestSqueezeCollapsedAxes: + + def test_no_collapsed_axes_is_a_passthrough(self): + grid = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 2.0, 4)] + values = np.ones((4, 3, 2)) + out_grid, out_values = squeeze_collapsed_axes(grid, values) + assert len(out_grid) == 2 + assert out_values.shape == (4, 3, 2) + + def test_drops_a_singleton_axis(self): + x = np.linspace(0.0, 1.0, 4) + y = np.array([0.5, 0.6]) # 1-cell axis (select()-ed) + z = np.linspace(-1.0, 1.0, 5) + values = np.zeros((3, 1, 4, 2)) + grid, out_values = squeeze_collapsed_axes([x, y, z], values) + assert len(grid) == 2 + assert out_values.shape == (3, 4, 2) + + def test_drops_multiple_singleton_axes(self): + x = np.linspace(0.0, 1.0, 4) + y = np.array([0.0]) + z = np.array([0.0]) + values = np.zeros((3, 1, 1, 2)) + grid, out_values = squeeze_collapsed_axes([x, y, z], values) + assert len(grid) == 1 + assert out_values.shape == (3, 2) + + def test_curvilinear_axis_is_averaged_not_indexed(self): + # A 2-D (curvilinear) coordinate array spanning both dims; dropping dim 1 + # (a singleton) should mean-reduce dim 1 out of the coordinate array too. + x2d = np.arange(12.0).reshape(4, 3) # (dim0=4 edges, dim1=3 edges) + y2d = np.arange(12.0).reshape(4, 3) * 2.0 + values = np.zeros((3, 1, 2)) # 3 cells in dim0, 1 cell in dim1 + grid, out_values = squeeze_collapsed_axes([x2d, y2d], values) + assert len(grid) == 1 + assert grid[0].shape == (4, ) + np.testing.assert_allclose(grid[0], np.mean(x2d, axis=1)) + assert out_values.shape == (3, 2) + + +# -------------------------------------------------------------------------- +# subplot_grid +# -------------------------------------------------------------------------- + + +class TestSubplotGrid: + + def test_perfect_square(self): + assert subplot_grid(4) == (2, 2) + + def test_single_panel(self): + assert subplot_grid(1) == (1, 1) + + def test_non_square_uses_near_square_layout(self): + rows, cols = subplot_grid(3) + assert rows * cols >= 3 + + def test_explicit_num_rows(self): + assert subplot_grid(6, num_rows=2) == (2, 3) + + def test_explicit_num_cols(self): + assert subplot_grid(6, num_cols=3) == (2, 3) + + def test_five_panels_layout(self): + rows, cols = subplot_grid(5) + assert rows * cols >= 5 + assert rows * cols <= 6 + + +# -------------------------------------------------------------------------- +# prep_plot_data +# -------------------------------------------------------------------------- + + +def _make_state(grid, values) -> GDataState: + d = GDataState() + d.push(grid, values) + return d + + +class TestPrepPlotData: + + def test_1d_basic(self): + grid = [np.linspace(0.0, 1.0, 9)] + values = np.ones((8, 1)) + panel = prep_plot_data(_make_state(grid, values)) + assert panel.num_dims == 1 + assert panel.num_comps == 1 + assert panel.xlabel == r"$z_0$" + assert panel.ylabel == "" + + def test_2d_basic(self): + grid = [np.linspace(0.0, 1.0, 5), np.linspace(0.0, 2.0, 6)] + values = np.ones((4, 5, 2)) + panel = prep_plot_data(_make_state(grid, values)) + assert panel.num_dims == 2 + assert panel.num_comps == 2 + assert panel.xlabel == r"$z_0$" + assert panel.ylabel == r"$z_1$" + + def test_squeezes_a_selected_axis(self): + x = np.linspace(0.0, 1.0, 4) + y = np.array([0.4, 0.6]) + values = np.zeros((3, 1, 2)) + panel = prep_plot_data(_make_state([x, y], values)) + assert panel.num_dims == 1 + assert panel.values.shape == (3, 2) + + def test_custom_xlabel_overrides_default(self): + grid = [np.linspace(0.0, 1.0, 5)] + values = np.ones((4, 1)) + panel = prep_plot_data(_make_state(grid, values), xlabel="time") + assert panel.xlabel == "time" + + def test_clabel_gets_zscale_annotation(self): + grid = [np.linspace(0.0, 1.0, 5)] + values = np.ones((4, 1)) + panel = prep_plot_data(_make_state(grid, values), clabel="n_e", zscale=2.0) + assert "2.00" in panel.clabel diff --git a/tests/test_render_pyvista.py b/tests/test_render_pyvista.py new file mode 100644 index 00000000..1630b873 --- /dev/null +++ b/tests/test_render_pyvista.py @@ -0,0 +1,254 @@ +"""Tests for postgkyl.render.pyvista -- 3-D volume/isosurface rendering. + +``pyvista`` is a hard dependency (pyproject.toml) but needs a working +(possibly software/off-screen) OpenGL context; every test here renders +off-screen (``no_show=True``) and is skipped cleanly if that context is not +available on the host, per the layer instructions. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +pv = pytest.importorskip("pyvista") + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.render.pyvista import pyvista + + +def _has_gl_context() -> bool: + # This host's VTK build reports "vtkXOpenGLRenderWindow" -- GLX-only, no + # OSMesa/EGL fallback -- so without a real or virtual (Xvfb) X server to + # connect to, VTK doesn't raise a catchable exception: it hits a fatal, + # unrecoverable "Fatal Python error: Aborted" that takes the whole pytest + # process down. Check for a display *before* touching pyvista/VTK at all, + # so a truly headless host skips instead of aborting the run. + if not os.environ.get("DISPLAY"): + return False + try: + pl = pv.Plotter(off_screen=True) + pl.add_mesh(pv.Sphere()) + pl.screenshot() + pl.close() + return True + except Exception: + return False + + +needs_gl = pytest.mark.skipif( + not _has_gl_context(), + reason="no working (off-screen) OpenGL context on this host") + + +def _volume(n=6) -> GDataState: + grid = [np.linspace(0.0, 1.0, n + 1) for _ in range(3)] + x, y, z = np.meshgrid(*[0.5 * (g[:-1] + g[1:]) for g in grid], indexing="ij") + values = (x + y + z)[..., np.newaxis] + d = GDataState() + d.push(grid, values) + return d + + +@needs_gl +class TestPyvista: + + def test_offscreen_volume_render_does_not_raise(self): + pyvista(_volume(), no_show=True, volume=True) + + def test_offscreen_contour_render_does_not_raise(self): + pyvista(_volume(), no_show=True, volume=False, contour_levels=4) + + def test_saves_a_png_screenshot(self, tmp_path): + out = tmp_path / "out.png" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert out.exists() + assert out.stat().st_size > 0 + + def test_saves_an_html_export(self, tmp_path): + # pyvista's HTML export needs the optional "trame" extra, not (only) a + # GL context; skip cleanly rather than mislabel it as a GL failure. + pytest.importorskip("trame_vtk") + out = tmp_path / "out.html" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert out.exists() + + def test_log_color_scale_does_not_raise(self): + pyvista(_volume(), no_show=True, is_log=True) + + def test_diverging_colormap_does_not_raise(self): + pyvista(_volume(), no_show=True, diverging=True) + + def test_clip_plane_does_not_raise(self): + pyvista(_volume(), no_show=True, mesh_clip_plane=True) + + def test_clip_plane_volume_mode_does_not_raise(self): + pyvista(_volume(), no_show=True, volume=True, mesh_clip_plane=True) + + def test_hide_axes_does_not_raise(self): + pyvista(_volume(), no_show=True, hide_axes=True) + + def test_cylindrical_to_cartesian_does_not_raise(self): + pyvista(_volume(), no_show=True, cylindrical_to_cartesian=True) + + def test_diverging_opacity_ramp_does_not_raise(self): + pyvista(_volume(), no_show=True, opacity="diverging") + + def test_named_theme_does_not_raise(self): + pyvista(_volume(), no_show=True, theme="document") + + def test_hide_zeros_hides_exact_zero_points(self): + d = _volume() + d.values[0, 0, 0, 0] = 0.0 + pyvista(d, no_show=True, hide_zeros=True) + + def test_mesh_slice_plane_contour_mode_does_not_raise(self): + pyvista(_volume(), no_show=True, volume=False, mesh_slice_plane=True) + + def test_mesh_slice_plane_volume_mode_does_not_raise(self): + pyvista(_volume(), no_show=True, volume=True, mesh_slice_plane=True) + + def test_volume_clip_plane_does_not_raise(self): + pyvista(_volume(), no_show=True, volume=True, volume_clip_plane=True) + + def test_saves_a_vector_graphic(self, tmp_path): + out = tmp_path / "out.svg" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert out.exists() + + def test_saves_a_gltf_export(self, tmp_path): + out = tmp_path / "out.gltf" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert out.exists() + + def test_saves_a_vtksz_export(self, tmp_path): + # Like .html, PyVista's .vtksz export needs the optional "trame" extra. + pytest.importorskip("trame_vtk") + out = tmp_path / "out.vtksz" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert out.exists() + + def test_no_title_omits_add_text(self): + pyvista(_volume(), no_show=True, title=None) + + def test_gl_context_errors_propagate_unwrapped(self): + # _require_gl_context only wraps *unexpected* exceptions from the render + # backend into a RuntimeError; a ValueError PyVista itself raises (e.g. + # an unknown theme name) should pass through as-is, not get relabeled as + # a GL-context failure. + with pytest.raises(ValueError, match="Theme"): + pyvista(_volume(), no_show=True, theme="bogus_theme_xyz") + + def test_spin_rotates_camera_and_stops_after_interaction(self, monkeypatch): + # The rotation timer/click-observer callbacks only ever run inside VTK's + # own interactive event loop, which off-screen tests never enter. Capture + # them by stubbing the registration calls, then invoke them directly to + # exercise the closures' logic (advance while idle, freeze on click). + from pyvista.plotting.render_window_interactor import RenderWindowInteractor + + captured = {} + monkeypatch.setattr( + pv.Plotter, "add_timer_event", + lambda self, max_steps, duration, callback: captured.setdefault( + "rotate", callback)) + + def _fake_add_observer(self, event, call, interactor_style_fallback=True): + if event == "LeftButtonPressEvent": + captured["click"] = call + + monkeypatch.setattr(RenderWindowInteractor, "add_observer", + _fake_add_observer) + + pyvista(_volume(), no_show=True, no_spin=False, volume=True) + + assert "rotate" in captured and "click" in captured + captured["rotate"](0) + captured["click"]() + captured["rotate"](0) # a no-op once "clicked": interacting freezes it + + def test_html_saveas_dispatches_to_export_html(self, monkeypatch, tmp_path): + # Exercise postgkyl's own saveas -> exporter dispatch without requiring + # the optional "trame_vtk" extra that pyvista's real HTML export needs. + called = {} + monkeypatch.setattr(pv.Plotter, "export_html", + lambda self, path: called.setdefault("path", path)) + out = tmp_path / "out.html" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert called["path"] == str(out) + + def test_vtksz_saveas_dispatches_to_export_vtksz(self, monkeypatch, tmp_path): + # Same as above, but for the optional "trame" extra .vtksz export needs. + called = {} + monkeypatch.setattr(pv.Plotter, "export_vtksz", + lambda self, path: called.setdefault("path", path)) + out = tmp_path / "out.vtksz" + pyvista(_volume(), no_show=True, saveas=str(out)) + assert called["path"] == str(out) + + def test_show_true_calls_plotter_show(self, monkeypatch): + # A real interactive .show() blocks waiting for the window to close; + # stub it out to exercise the no_show=False branch without hanging the test. + calls = [] + monkeypatch.setattr(pv.Plotter, "show", + lambda self, *a, **k: calls.append(True)) + pyvista(_volume(), no_show=False, volume=True) + assert calls == [True] + + def test_show_bounds_axes_ranges_reflect_scale_and_shift(self, monkeypatch): + # The mesh itself is always normalized to +/-aspect_ratio (PyVista + # handles non-integer axis extents poorly), so axes_ranges is the only + # thing that can carry the user's requested xscale/yscale/zscale and + # xshift/yshift/zshift into the displayed tick labels -- see C1. + captured = {} + original_show_bounds = pv.Plotter.show_bounds + + def _spy(self, **kwargs): + captured.update(kwargs) + return original_show_bounds(self, **kwargs) + + monkeypatch.setattr(pv.Plotter, "show_bounds", _spy) + + grid = [np.linspace(0.0, 1.0, 7) for _ in range(3)] + centers = 0.5 * (grid[0][:-1] + grid[0][1:]) + xmin, xmax = float(centers.min()), float(centers.max()) + x, y, z = np.meshgrid(centers, centers, centers, indexing="ij") + values = (x + y + z)[..., np.newaxis] + d = GDataState() + d.push(grid, values) + + pyvista(d, + no_show=True, + volume=True, + xscale=2.0, + xshift=1.0, + yscale=3.0, + yshift=0.5, + zscale=4.0, + zshift=1.0) + + assert "axes_ranges" in captured + axes_ranges = captured["axes_ranges"] + # Volume mode with the default aspect_ratio=(1,1,1) and no downsampling + # builds a mesh spanning exactly [-1, 1] per axis, so pv_bounds.*_min/ + # *_max are -1/+1 and axes_ranges reduces to (val + shift) * scale. + np.testing.assert_allclose(axes_ranges[0], (xmin + 1.0) * 2.0, atol=1e-9) + np.testing.assert_allclose(axes_ranges[1], (xmax + 1.0) * 2.0, atol=1e-9) + np.testing.assert_allclose(axes_ranges[2], (xmin + 0.5) * 3.0, atol=1e-9) + np.testing.assert_allclose(axes_ranges[3], (xmax + 0.5) * 3.0, atol=1e-9) + np.testing.assert_allclose(axes_ranges[4], (xmin + 1.0) * 4.0, atol=1e-9) + np.testing.assert_allclose(axes_ranges[5], (xmax + 1.0) * 4.0, atol=1e-9) + + +class TestPyvistaValidation: + + def test_non_3d_dataset_raises(self): + d = GDataState() + d.push([np.linspace(0.0, 1.0, 5)], np.ones((4, 1))) + with pytest.raises(ValueError, match="3D"): + pyvista(d, no_show=True) + + def test_unsupported_saveas_extension_raises(self): + with pytest.raises(ValueError, match="Unsupported"): + pyvista(_volume(), no_show=True, saveas="out.bogus") diff --git a/tests/test_render_style.py b/tests/test_render_style.py new file mode 100644 index 00000000..cfd26760 --- /dev/null +++ b/tests/test_render_style.py @@ -0,0 +1,49 @@ +"""Tests for postgkyl.render.style -- apply_style.""" + +from __future__ import annotations + +import matplotlib + +matplotlib.use("Agg") +import matplotlib as mpl +import pytest + +from postgkyl.render.style import DEFAULT_STYLE, apply_style + + +@pytest.fixture(autouse=True) +def _restore_rcparams(): + with mpl.rc_context(): + yield + + +class TestApplyStyle: + + def test_default_applies_packaged_postgkyl_style(self): + apply_style() + assert mpl.rcParams["image.cmap"] == "inferno" + assert mpl.rcParams["image.origin"] == "lower" + + def test_named_postgkyl_style_matches_default(self): + apply_style(DEFAULT_STYLE) + assert mpl.rcParams["image.cmap"] == "inferno" + + def test_cycler_line_is_parsed_by_matplotlib(self): + apply_style() + cycle = list(mpl.rcParams["axes.prop_cycle"]) + assert len(cycle) == 7 + + def test_matplotlib_named_style_is_forwarded(self): + apply_style("default") + # "default" resets to Matplotlib's own baseline cmap. + assert mpl.rcParams["image.cmap"] == "viridis" + + def test_arbitrary_mplstyle_path_is_applied(self, tmp_path): + style_file = tmp_path / "custom.mplstyle" + style_file.write_text("image.cmap: plasma\n") + apply_style(str(style_file)) + assert mpl.rcParams["image.cmap"] == "plasma" + + def test_unknown_style_name_raises(self): + with pytest.raises(OSError): + apply_style("this-style-does-not-exist") diff --git a/tests/test_select.py b/tests/test_select.py deleted file mode 100644 index 48986226..00000000 --- a/tests/test_select.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Postgkyl module for testing data loading.""" -import numpy as np -import os - -import postgkyl as pg - - -class TestSelect: - """Test Gkeyll's select commands.""" - dir_path = f"{os.path.dirname(__file__)}/test_data" - - def test_integer(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - grid, data = pg.data.select(data, z0=1, z1="5:8") - np.testing.assert_array_equal(grid[0], [1.375, 1.75 ]) - np.testing.assert_array_almost_equal(grid[1], [3.926991, 4.712389, 5.497787, 6.283185]) - np.testing.assert_array_equal(data.shape, (1, 3, 4)) - - def test_float(self): - data = pg.GData(f"{self.dir_path:s}/shock-f-ser-p1.gkyl") - grid, data = pg.data.select(data, z0=0.5) - np.testing.assert_array_equal(grid[0], [1.0, 1.375]) diff --git a/tests/test_version_report.py b/tests/test_version_report.py new file mode 100644 index 00000000..538145e5 --- /dev/null +++ b/tests/test_version_report.py @@ -0,0 +1,85 @@ +"""Deterministic coverage for version-report fallback paths.""" + +from __future__ import annotations + +import importlib.metadata +from importlib import import_module +import subprocess + +import pytest + +version = import_module("postgkyl._version") + +pytestmark = pytest.mark.compatibility + + +def test_git_returns_none_outside_a_checkout(tmp_path): + assert version._git(tmp_path, "status") is None + + +@pytest.mark.parametrize( + "error", [OSError("missing git"), + subprocess.CalledProcessError(1, "git")]) +def test_git_converts_process_failures_to_missing(monkeypatch, tmp_path, error): + (tmp_path / ".git").mkdir() + + def fail(*_args, **_kwargs): + raise error + + monkeypatch.setattr(version.subprocess, "run", fail) + assert version._git(tmp_path, "status") is None + + +@pytest.mark.parametrize( + ("build", "expected"), + [ + (None, "unknown (not a git checkout)"), + ({ + "postgkyl_build_commit": "unknown" + }, "unknown (not a git checkout)"), + ({ + "postgkyl_build_commit": "1234567890abcdef" + }, "1234567890ab (baked at build time, not a git checkout)"), + ], +) +def test_postgkyl_commit_uses_baked_fallback(monkeypatch, build, expected): + monkeypatch.setattr(version, "_git", lambda *_args: None) + monkeypatch.setattr(version.gpython, "build_info", lambda: build) + assert version._postgkyl_commit() == expected + + +def test_postgkyl_commit_marks_a_dirty_checkout(monkeypatch): + answers = iter(["abcdef123456", " M changed.py"]) + monkeypatch.setattr(version, "_git", lambda *_args: next(answers)) + assert version._postgkyl_commit() == "abcdef123456-dirty" + + +def test_gkeyll_info_reports_an_unbuilt_bridge(monkeypatch): + monkeypatch.setattr(version.gpython, "build_info", lambda: None) + assert "not built" in version._gkeyll_info() + + +def test_dependency_versions_omits_missing_distributions(monkeypatch): + + def distribution_version(name): + if name == "scipy": + raise importlib.metadata.PackageNotFoundError(name) + return "1.2.3" + + monkeypatch.setattr(version.importlib.metadata, "version", + distribution_version) + report = version._dependency_versions() + assert "numpy 1.2.3" in report + assert "scipy" not in report + + +def test_version_report_without_build_metadata(monkeypatch): + monkeypatch.setattr(version.gpython, "build_info", lambda: None) + monkeypatch.setattr(version.gpython, "available", lambda: False) + monkeypatch.setattr(version, "_postgkyl_commit", lambda: "unknown") + monkeypatch.setattr(version, "_gkeyll_info", lambda: "not built") + monkeypatch.setattr(version, "_dependency_versions", lambda: "numpy 1") + report = version.version_report("2.0") + assert "pgkyl, version 2.0" in report + assert "gpython bridge: unavailable" in report + assert "ARCH_FLAGS" not in report