diff --git a/.agents/skills/develop-test-driven/SKILL.md b/.agents/skills/develop-test-driven/SKILL.md new file mode 100644 index 000000000..e1ac675d4 --- /dev/null +++ b/.agents/skills/develop-test-driven/SKILL.md @@ -0,0 +1,231 @@ +--- +name: develop-test-driven +description: Enforce real red-green-refactor implementation instead of code-first changes. Use this whenever a task will change executable behavior and the repo has a practical automated test seam, whether directly from the user or as a delegated slice inside `execute-plan`. +compatibility: opencode +--- + +# Develop test driven + +Use this skill when implementation should be led by tests instead of by coding first and adding coverage afterward. The goal is to change behavior with evidence, not to write code that merely seems right. + +This skill is for active implementation work. It is most useful for features, bug fixes, and refactors where the correct behavior can be expressed in automated tests. It works well as a focused sub-skill inside `execute-plan`: `execute-plan` owns sequencing, delegation, work card creation, and integrated verification handoff; this skill owns the red-green-refactor loop for the implementation slice. + +When the slice is mainly about an interface boundary such as web UI, CLI, TUI, API contracts, emails, or notifications, pair this skill with `implement-interface` rather than forcing TDD to also carry all of the interface-fidelity guidance. + +If you were invoked from `execute-plan`, treat the feature card and current task as your boundary. Implement only that slice, return evidence, and let `execute-plan` decide broader sequencing and downstream handoffs. + +## Owns + +- One behavior-changing implementation slice at a time +- The red-green-refactor loop for that slice +- Producing concrete evidence: the failing test or automated check, the minimal code change, and the passing verification + +## Does Not Own + +- moving feature cards between lanes +- deciding overall task sequencing across a full feature plan +- product discovery, feature design, or implementation planning +- broad final review or documentation closure for the feature + +## When To Use + +Use this skill when the user or coordinating skill: + +- asks for a feature or bug fix and the repo has a test harness +- wants true TDD or says to write tests first +- needs a regression test for a reported failure before changing code +- is working in an area where behavior drift would be expensive +- is executing a task from `execute-plan` that changes executable behavior + +Do not use this skill when: + +- the change is purely docs, comments, planning artifacts, or static configuration with no meaningful behavior to test +- the repo genuinely has no practical automated test seam for the requested change +- the task is still discovery or planning and implementation should not start yet +- the main work is lane movement or feature-card maintenance rather than behavior change + +If there is no good automated seam yet, first find the smallest useful seam. Do not use that as an excuse to skip testing entirely. + +## Core Rule + +Do not write or keep production code for a new behavior change until a test or other automated check demonstrates the missing behavior first. + +That means: + +- write the test before the implementation change +- run it and confirm it fails for the expected reason +- make the smallest implementation change that can make it pass +- rerun the focused check and confirm the behavior changed +- refactor only after the tests are green + +If code was already written for the requested behavior before the failing test exists, do not quietly continue from there. Either delete the speculative code or ignore it and re-derive the implementation from the test. + +Speed is not an exception. When the user wants a quick fix, make the cycle smaller and more targeted, not looser. + +## Workflow + +### 1. Start From The Intended Behavior And The Current Task Boundary + +Before touching code, state what behavior is changing. + +Capture: + +- the user-visible or externally observable behavior +- the seam where that behavior can be tested +- the smallest first test that would prove the behavior is currently missing +- any task constraints supplied by `execute-plan` or the feature card + +Prefer a test that expresses the desired behavior through the real public interface. Default away from mocks when real interactions are practical. A mock-heavy test that only proves the mock was called is weaker than a small real test at the correct boundary. + +If a coordinating skill supplied task-specific verification, satisfy that verification without expanding into adjacent plan work. + +### 2. Write The Smallest Failing Test + +Add one focused test for one behavior. + +Good tests usually: + +- name the behavior clearly +- assert one meaningful outcome +- use real code paths when practical +- stay small enough that the reason for failure is obvious + +Avoid combining multiple behaviors into one broad test unless the repo's existing patterns clearly prefer that shape. + +### 3. Verify Red Before Implementing + +Run the most targeted command that proves the new test fails. + +Confirm all of these: + +- the test actually runs +- it fails, not passes +- it fails for the expected reason +- the failure points to missing or incorrect behavior, not a typo or broken test setup + +If the test passes immediately, you have not proven the change. Tighten the test or choose a better seam. + +If the test fails for the wrong reason, fix the test first. Do not start implementing against a broken signal. + +### 4. Make The Smallest Green Change + +Only after the failing signal is real, change the implementation. + +Aim for the minimum code needed to satisfy the current test. Prefer a direct, boring change over early abstraction. Do not bundle unrelated cleanup, extra options, adjacent feature work, or feature-card edits into the same step. + +When fixing a bug, the regression test is the first deliverable. The code change is second. + +### 5. Verify Green Immediately + +Rerun the focused test or check that failed in the red step. + +Then confirm nearby regression safety with whatever additional verification fits the repo and change risk, such as: + +- the relevant test file or package +- related integration tests +- lint, typecheck, or build when they are meaningful safeguards + +Do not claim success because the code looks correct. Claim success when the targeted failing signal turned green and the surrounding checks still hold. + +### 6. Refactor Without Changing Behavior + +Once green, improve the code only if the cleanup helps and the tests still guard the behavior. + +Good refactors here include: + +- removing duplication introduced during the green step +- clarifying names or structure +- extracting a helper that the passing tests now protect + +Do not sneak in extra behavior during refactor. If behavior must change again, start another red-green cycle. + +### 7. Return Evidence, Then Repeat If Needed + +For multi-part work, repeat with the next smallest missing behavior instead of writing a large block of implementation and backfilling tests later. + +Each cycle should leave behind: + +- one clearer expression of expected behavior +- one passing proof that the behavior now exists +- no ambiguity about what changed + +If you are working under `execute-plan`, report back in a way the coordinator can use directly: + +- the behavior covered +- the new or updated test +- the command that proved red +- the minimal implementation change for green +- the command that proved green +- any remaining risk or next likely cycle + +Do not move the feature card between lanes from this skill unless the coordinating instructions explicitly made that the current task, which is unusual. + +## TDD Heuristics + +Prefer: + +- user-facing or contract-facing tests over implementation-detail tests +- focused commands that validate one new behavior quickly +- minimal implementation before abstraction +- regression tests for every confirmed bug fix +- small repeated cycles over one giant test file rewrite + +Flag: + +- implementation code appearing before any new failing test +- a new test that passes on the first run without proving anything +- tests that mostly assert mocks, stubs, or internals instead of behavior +- broad code changes justified by one narrow failing test +- claims that manual checking is enough for a behavior that can be automated + +## When The Test Seam Is Hard To Reach + +Treat a hard-to-test change as design feedback. + +First ask: + +- is there a smaller public seam to test? +- can the behavior be moved behind a simpler interface? +- can an integration test cover this more honestly than deep mocking? + +If the only path forward is limited, be explicit about the compromise and still preserve the red-green order. A weaker but honest automated check is usually better than code-first work with no durable proof at all. + +If you are operating under `execute-plan` and no practical automated seam exists, return that constraint clearly so the coordinator can decide whether the task needs replanning, a documented exception, or a different verification approach. + +## Resist Common Drift + +Push back on these failure patterns: + +- "fix it first and add tests after" +- "just patch the function quickly" +- "manual verification is enough for this one" +- "the change is too small to bother testing" + +Do not argue abstractly. Translate the pushback into a smaller TDD step: name the narrow failing test, run it, make it pass, then continue. + +## Checkpoint Reporting + +At meaningful points, report progress in the language of the cycle: + +- the behavior under test +- the current feature-task boundary if one was supplied +- the new failing test or check +- the command that proved red +- the implementation change made for green +- the command that proved green +- any remaining behaviors that still need another cycle + +Keep these updates short. The point is to show evidence, not to narrate every keystroke. + +## Quality Bar + +Before finishing, check that: + +- every intended behavior change in the current slice is covered by a new or updated automated check +- each important new test was observed failing before the corresponding implementation change +- the final implementation is no broader than the tested scope requires +- bug fixes include a regression test that would have caught the original failure +- the relevant verification commands were actually run and their outcomes are known +- any coordination context from `execute-plan` was honored without drifting into unrelated work + +If those are not true, the work is not done yet. Keep cycling until the evidence matches the claim. diff --git a/.github/skills/port-quantity/SKILL.md b/.github/skills/port-quantity/SKILL.md new file mode 100644 index 000000000..a143f46d6 --- /dev/null +++ b/.github/skills/port-quantity/SKILL.md @@ -0,0 +1,557 @@ +--- +name: port-quantity +description: "Port a py4vasp quantity from the inheritance-based Refinery architecture to the new Dispatcher/Impl architecture. USE WHEN: migrating an existing quantity class, porting tests, or adding a new quantity. Triggers: 'port quantity', 'migrate to new architecture', 'convert to composition', 'new architecture', 'refactor quantity'." +--- + +# Port a py4vasp Quantity to the Dispatcher/Handler Architecture + +Port an existing `Refinery`-based quantity to the new architecture described in `docs/architecture/calculation.rst`. + +## Core Design Contract + +Each quantity is split into two classes: + +- **Dispatcher** (public, e.g. ``Band``) — user-facing, attached to ``Calculation``. Owns the ``Source``, calls standalone dispatch functions that parse selections, open data, construct the Handler, call methods, and merge results. +- **Handler** (private, e.g. ``BandHandler``) — constructed via ``from_data(raw)``. Works with exactly one raw data object. Contains all transform logic. Primary unit-testing target. + +The dispatcher does **not** have ``from_data``. That lives exclusively on the Handler. + +### Selection dispatch + +Selection dispatch is handled by standalone functions named after their merge +strategy. The dispatcher just calls the appropriate one: + +```python +def merge_default(source, quantity_name, selection, handler_factory, method, *args, **kwargs): + """Use this if no specific merge strategy is required. This will merge the results + into a dict by selection name if more than one result is returned.""" + ... + +def merge_graphs(source, quantity_name, selection, handler_factory, method, *args, **kwargs): + """Overlay Graph results into a single figure.""" + ... + +def merge_strings(source, quantity_name, selection, handler_factory, method, *args, **kwargs): + """Return a single string by concatenating results for all selections.""" + ... +``` + +**Choosing the merge strategy:** + +- `merge_default` — **the default** for every method. Use this almost always. +- `merge_graphs` — for methods that return a `Graph` (typically `plot`). +- `merge_strings` — for methods that return a string (e.g. `__str__`). + +All call the inner `_dispatch` which: +1. Parses `selection` via `_parse_selections` → list of `SelectionContext(selection_name, remaining_selection)`. +2. For each context, calls `source.access(quantity_name, selection=ctx.selection_name)`. +3. Inside the context, calls `handler_factory(raw)` then `method(handler, *args, **kwargs)`. +4. Collects results as `{selection_name: result}`. + +```python +class SelectionContext(typing.NamedTuple): + selection_name: str | None + remaining_selection: str | None +``` + +### Handler pattern + +Type hints on `_raw_band` and in `from_data` are **mandatory** — never omit them. +They document which raw dataclass the Handler works with and are essential for +readability and static analysis. + +```python +from py4vasp import raw + +class BandHandler: + def __init__(self, raw_band: raw.Band): + self._raw_band = raw_band + + @classmethod + def from_data(cls, raw_band: raw.Band) -> "BandHandler": + return cls(raw_band) + + def read(self) -> dict: + return { + "eigenvalues": np.array(self._raw_band.eigenvalues) - self._raw_band.fermi_energy, + ... + } + + def plot(self) -> Graph: + ... +``` + +### Dispatcher pattern + +```python +@quantity("band") +class Band: + def __init__(self, source: Source, quantity_name: str = "band"): + self._source = source + self._quantity_name = quantity_name + + def read(self, selection: str | None = None) -> dict: + return merge_default( + self._source, self._quantity_name, selection, + BandHandler.from_data, BandHandler.read, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Public alias for read(). Part of the public interface — never deprecate.""" + return self.read(selection=selection) + + def plot(self, selection: str | None = None) -> Graph: + return merge_graphs( + self._source, self._quantity_name, selection, + BandHandler.from_data, BandHandler.plot, + ) +``` + +**Important:** Every dispatcher method that has `@base.data_access` in the old +Refinery code has `selection` injected into it at runtime. After porting, **all** +dispatcher methods must declare `selection: str | None = None` in their signature, +even if the Refinery method had no explicit `selection` parameter. + +**`to_dict` is part of the public interface and must never be deprecated or +removed.** In the new architecture, `to_dict` is a thin alias that delegates to +`read()`. Tests for `to_dict` should verify it returns the same result as `read()`: + +```python +# Tests for to_dict +def test_to_dict_matches_read(raw_bandgap): + handler = BandgapHandler.from_data(raw_bandgap) + assert handler.to_dict() == handler.read() # or via dispatcher: + +def test_dispatcher_to_dict_matches_read(raw_bandgap): + source = DataSource(raw_bandgap) + dispatcher = Bandgap(source=source, quantity_name="bandgap") + assert dispatcher.to_dict() == dispatcher.read() +``` + +Extra arguments from the dispatcher method are forwarded: + +```python +def plot(self, selection=None, fermi_energy=None): + return merge_graphs( + self._source, self._quantity_name, selection, + BandHandler.from_data, BandHandler.plot, + fermi_energy=fermi_energy, + ) +``` + +--- + +## Migration Procedure + +### 1 — Identify the raw dataclass + +Open `src/py4vasp/_raw/data.py`. Find the dataclass matching this quantity (CamelCase of the quantity name). This becomes the type for the Handler's `_raw_` attribute. Note all fields and their types. + +Example for `bandgap`: +```python +@dataclasses.dataclass +class Bandgap: + labels: VaspData + values: VaspData +``` + +### 2 — Create the Handler class + +Create a private `Handler` class. It takes raw data in its constructor and has a `from_data` classmethod. Move all transform logic here. + +Type hints on `_raw_` and in `from_data` are **mandatory**. Always use the exact +raw dataclass type from `py4vasp.raw` defined in `py4vasp/_raw/data.py`. + +```python +# Before (on the Refinery) +class Bandgap(slice_.Mixin, base.Refinery, graph.Mixin): + _raw_data: raw_data.Bandgap + + @base.data_access + def to_dict(self): + return { + **self._gap_dict("fundamental"), + ... + "fermi_energy": self._get("Fermi energy", component=0), + } + +# After (Handler) +class BandgapHandler: + def __init__(self, raw_bandgap: raw.Bandgap, steps=None): + self._raw_bandgap = raw_bandgap + self._steps = steps + + @classmethod + def from_data(cls, raw_bandgap: raw.Bandgap, steps=None) -> "BandgapHandler": + return cls(raw_bandgap, steps=steps) + + def read(self) -> dict: + return { + **self._gap_dict("fundamental"), + ... + "fermi_energy": self._get("Fermi energy", component=0), + } +``` + +Key changes: +- Replace `self._raw_data` with `self._raw_` everywhere in the Handler. +- Remove `@base.data_access` decorators — Handler methods are plain methods. +- The Handler never touches `Source` or selection dispatch. + +### 3 — Create the Dispatcher class + +Create the public class with the `@quantity()` decorator. It owns the `Source` and delegates to dispatch functions. + +```python +@quantity("bandgap") +class Bandgap(graph.Mixin): + def __init__(self, source: Source, quantity_name: str = "bandgap", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + def __getitem__(self, steps) -> Bandgap: + return Bandgap(self._source, self._quantity_name, steps=steps) + + def _handler_factory(self, raw): + return BandgapHandler.from_data(raw, steps=self._steps) + + def read(self, selection: str | None = None) -> dict: + return merge_default( + self._source, self._quantity_name, selection, + self._handler_factory, BandgapHandler.read, + ) + + def to_dict(self, selection: str | None = None) -> dict: + return self.read(selection=selection) + + def plot(self, selection: str | None = None) -> Graph: + return merge_graphs( + self._source, self._quantity_name, selection, + self._handler_factory, BandgapHandler.plot, + ) +``` + +For step-indexed quantities, use `self._handler_factory` as a bound method that captures `self._steps` via the partial pattern shown above. + +### 3.5 — Port docstrings to the Dispatcher + +Every public method docstring from the Refinery **must** be carried over to the +corresponding Dispatcher method. Do not leave bare one-liners in place of the +original prose. + +The `@documentation.format` decorator and `slice_.examples` helper work exactly +as in the old architecture — keep using them on the Dispatcher: + +```python +from py4vasp._calculation import slice_ +from py4vasp._util import documentation + +@quantity("bandgap") +@documentation.format(examples=slice_.examples("bandgap")) +class Bandgap(graph.Mixin): + """...original class docstring... + + {examples} + """ + + @documentation.format(examples=slice_.examples("bandgap", "read")) + def read(self, selection: str | None = None) -> dict: + """Read the bandgap data from a VASP relaxation or MD trajectory. + + Returns + ------- + dict + Contains ... + + {examples} + """ + + @documentation.format(examples=slice_.examples("bandgap", "to_dict")) + def to_dict(self, selection: str | None = None) -> dict: + """Read the bandgap data from a VASP relaxation or MD trajectory. + + Convenient alias for :py:meth:`read`. Check that method for examples + and optional arguments. + + {examples} + """ +``` + +**Rules:** +- Apply `@documentation.format(examples=slice_.examples("quantity", "method"))` to every + method that had it on the Refinery. +- Apply `@documentation.format(examples=slice_.examples("quantity"))` to the class itself. +- `to_dict` on the Dispatcher is an alias — give it a short docstring noting the alias + relationship, but keep the `@documentation.format` decorator so the examples are present. +- Do **not** port `@documentation.format` to the Handler — its methods are internal + and do not need user-facing examples. + +### 4 — Move private helpers to the Handler + +Private helpers (`_gap`, `_get`, `_kpoint`, `_spin_polarized`, etc.) that read `self._raw_data` move to the Handler and read `self._raw_` instead. + +```python +# Before (Refinery) +def _spin_polarized(self): + return self._raw_data.values.shape[1] == 3 + +# After (Handler) +def _spin_polarized(self): + return self._raw_bandgap.values.shape[1] == 3 +``` + +### 5 — Handle selection forwarding + +`_dispatch` does **not** automatically forward `remaining_selection` to the +Handler. `SelectionContext.remaining_selection` is used only internally to +identify which source to open. + +If a Handler method needs to parse a selection string further (e.g. filtering +spin channels or orbital projections), pass the full `selection` string +explicitly as a keyword argument from the dispatcher: + +```python +# Dispatcher — forward selection explicitly when the handler needs it +def plot(self, selection=None): + return merge_graphs( + self._source, self._quantity_name, selection, + self._handler_factory, BandgapHandler.plot, + selection=selection, + ) + +# Handler — receives the full selection string and parses it itself +class BandgapHandler: + def plot(self, selection=None) -> Graph: + tree = select.Tree.from_selection(selection) + ... +``` + +For quantities whose Handler methods do not need further selection parsing, +omit the `selection=selection` kwarg entirely. + +### 6 — Handle composition with other quantities + +Use the other Impl's `from_data` directly — no Source needed: + +```python +class DensityHandler: + def read(self) -> dict: + structure = _StructureHandler.from_data(self._raw_density.structure) + return { + "charge": np.array(self._raw_density.charge), + "structure": structure.read(), + } +``` + +### 7 — Step-indexed quantities + +Steps live on the dispatcher and are passed to the Handler via the factory: + +```python +# Dispatcher +def __getitem__(self, steps) -> Structure: + return Structure(self._source, self._quantity_name, steps=steps) + +def _handler_factory(self, raw): + return StructureHandler.from_data(raw, steps=self._steps) +``` + +The Handler applies `slice_steps` explicitly: + +```python +from py4vasp._calculation.dispatch import slice_steps + +class StructureHandler: + def __init__(self, raw_structure: raw.Structure, steps=None): + self._raw_structure = raw_structure + self._steps = steps + + def read(self) -> dict: + return { + "lattice_vectors": slice_steps( + np.array(self._raw_structure.lattice_vectors), self._steps, default_ndim=2 + ), + } +``` + +`slice_steps(data, steps, default_ndim)` rules: +- `steps=None` → last step (default) +- `steps=3` → single step +- `steps=slice(1, 8)` → range +- `data.ndim <= default_ndim` → no step axis, return unchanged + +### 8 — Port `__str__` and display methods + +Move the string formatting logic to the Handler. The dispatcher calls it through dispatch: + +```python +# Handler +class BandgapHandler: + def __str__(self): + template = """...""" + return template.format(...) + +# Dispatcher +class Bandgap: + def __str__(self): + return merge_strings( + self._source, self._quantity_name, None, + self._handler_factory, BandgapHandler.__str__, + ) +``` + +### 9 — Port `_to_database` + +Database serialization lives **exclusively on the Handler** as a public method. +The Dispatcher has **no** database method — no `to_database`, no `_read_to_database`. + +```python +# Handler only — public, no leading underscore +class BandgapHandler: + def to_database(self) -> dict: + bandgap_dict = {...} + return {"bandgap": Bandgap_DB(**final_dict)} + +# Dispatcher — nothing. Do NOT add a database method here. +``` + +The method must return a `dict` with a single key (the quantity name) mapping to +a dataclass instance from `_raw/data_db.py`. Database consumers call +`Handler.from_data(raw).to_database()` directly. + +Tests for `to_database` use a `BandgapHandler` instance, not the dispatcher: + +```python +@pytest.fixture +def bandgap_handler(raw_data): + return setup_bandgap_handler(raw_data.bandgap("default")) + +def test_to_database_default(bandgap_handler, Assert): + _check_to_database(bandgap_handler, Assert) + +def _check_to_database(handler, Assert): + db_data = handler.to_database() # call on handler, not dispatcher + ... +``` + +### 10 — Port the tests + +**Never remove an existing test.** If a test cannot work yet (e.g. because the +dispatcher infrastructure isn't fully wired or factory methods changed), mark it +with `@pytest.mark.skip(reason="...")` so it remains visible and will be +re-enabled later. + +**`to_dict` tests:** `to_dict` is a public method — do not skip or delete its +tests. Instead, restructure them so they verify that `to_dict` and `read` +return the same result: + +```python +@pytest.mark.skip(reason="Dispatcher not yet wired to Calculation") +def test_factory_methods(raw_data, check_factory_methods): + ... +``` + +```python +# to_dict test: verify it equals read() +def test_to_dict_matches_read(raw_bandgap): + handler = BandgapHandler.from_data(raw_bandgap) + assert handler.to_dict() == handler.read() +``` + +Tests split into two categories: + +**Unit tests (Handler directly, no I/O):** + +```python +def test_bandgap_read(): + raw = raw_data.Bandgap(labels=..., values=...) + handler = BandgapHandler.from_data(raw) + result = handler.read() + assert ... + +def test_bandgap_step_selection(): + raw = raw_data.Bandgap(labels=..., values=...) + handler = BandgapHandler.from_data(raw, steps=3) + result = handler.read() + assert ... +``` + +**Integration tests (full pipeline via DictSource):** + +```python +def test_bandgap_via_calculation(): + source = DictSource({"bandgap": raw.Bandgap(...)}) + calc = Calculation(source=source) + result = calc.bandgap.read() + assert ... +``` + +The existing `from_data` pattern in tests like: +```python +bandgap = Bandgap.from_data(raw_bandgap) +``` +should be migrated to: +```python +handler = BandgapHandler.from_data(raw_bandgap) +``` + +Or for testing the dispatcher with the full dispatch pipeline: +```python +source = DataSource(raw_bandgap) +bandgap = Bandgap(source=source, quantity_name="bandgap") +``` + +### 11 — Remove from `QUANTITIES` / `GROUPS` + +In `src/py4vasp/_calculation/__init__.py`, remove the quantity's string entry from `QUANTITIES` (or `GROUPS`). The `@quantity()` decorator handles registration automatically. + +### 12 — Verify + +```bash +pytest tests/calculation/test_{name}.py -v # quantity-level tests +pytest tests/ -x # full suite +``` + +--- + +## Checklist + +For each quantity being ported: + +- [ ] Raw dataclass type identified in `_raw/data.py` +- [ ] Impl class created: `Handler` with `from_data(raw, steps=None)` +- [ ] All transform logic moved from Refinery to Handler +- [ ] `self._raw_data.x` → `self._raw_.x` in Handler +- [ ] `@base.data_access` decorators removed +- [ ] Dispatcher class created with `@quantity("name")` decorator +- [ ] Docstrings ported from Refinery to Dispatcher; `@documentation.format` + `slice_.examples` applied where the Refinery had them +- [ ] All dispatcher methods have `selection: str | None = None` parameter +- [ ] Dispatcher calls `merge_default` (default) or `merge_graphs` (for Graph) — no lambdas; `merge_strings` (for strings) +- [ ] `to_dict` kept in both Handler and dispatcher; dispatcher delegates to `read()`; not deprecated +- [ ] Type hints on `_raw_` attribute and `from_data` classmethod (mandatory, never omit) +- [ ] Handler method passed as unbound reference: `BandHandler.read` +- [ ] Extra args forwarded via `*args, **kwargs` +- [ ] Step indexing: `__getitem__` on dispatcher, `_handler_factory` captures steps +- [ ] Composition: other Handler's `from_data` called directly +- [ ] `__str__` / `_repr_pretty_` ported through dispatch +- [ ] `Handler.to_database()` implemented (public, no underscore); Dispatcher has no database method +- [ ] Small mixins kept (e.g. `graph.Mixin`) +- [ ] `base.Refinery` and `slice_.Mixin` removed from inheritance +- [ ] Tests split: Handler unit tests + dispatcher integration tests +- [ ] Tests for `to_dict` restructured to verify it matches `read()` — not skipped +- [ ] Tests never removed — non-working tests marked `@pytest.mark.skip(reason="...")` +- [ ] Removed from `QUANTITIES`/`GROUPS` in `__init__.py` +- [ ] All non-skipped tests pass + +--- + +## Reference Files + +| File | Purpose | +|------|---------| +| `docs/architecture/calculation.rst` | Full architecture description | +| `src/py4vasp/_raw/data.py` | Raw dataclass definitions | +| `src/py4vasp/_raw/definition.py` | Schema (sources per quantity) | +| `src/py4vasp/_calculation/bandgap.py` | Reference: existing Refinery quantity | +| `tests/calculation/test_bandgap.py` | Reference: existing test structure | diff --git a/TODOs.txt b/TODOs.txt new file mode 100644 index 000000000..13fd0478d --- /dev/null +++ b/TODOs.txt @@ -0,0 +1,10 @@ +FAILED tests/analysis/test_mlff.py::test_read_inputs_from_path - py4vasp.exception._Py4VaspInternalError: The selected dimension 1 is outside of the dimension of the data range(0, 1). +FAILED tests/raw/test_data.py::test_operators - hypothesis.errors.FailedHealthCheck: Input generation is slow: Hypothesis only generated 2 valid inputs after 2.88 seconds. +FAILED tests/scripts/test_error_analysis.py::test_error_analysis - assert 1 == 0 +FAILED tests/util/test_database.py::test_demo_db[None-5] - assert 4 > 5 +FAILED tests/util/test_database.py::test_demo_db[collinear-1] - assert False +FAILED tests/util/test_database.py::test_demo_db[noncollinear-5] - assert 2 > 5 +FAILED tests/util/test_database.py::test_demo_db[spin_texture-2] - assert 1 > 2 +FAILED tests/util/test_database.py::test_demo_db_with_tags[None] - assert False +FAILED tests/util/test_database.py::test_demo_db_with_tags[test] - assert False +FAILED tests/util/test_database.py::test_demo_db_with_tags[tags2] - assert False \ No newline at end of file diff --git a/plan.md b/plan.md new file mode 100644 index 000000000..0bb86ca90 --- /dev/null +++ b/plan.md @@ -0,0 +1,203 @@ +# Port All Quantities to Dispatcher/Handler Architecture + +Reference: `docs/architecture/calculation.rst`, skill in `.github/skills/port-quantity/SKILL.md` +Reference implementation (already ported): `src/py4vasp/_calculation/bandgap.py` + `tests/calculation/test_bandgap.py` + +## Status + +- [x] `bandgap` — reference port, complete + +--- + +## Porting Order + +Quantities are ordered by dependency. Port prerequisites before dependents. + +--- + +### Group 1 — Simple (no composition, no step-indexing) + +These are self-contained and have no internal dependencies on other quantities. + +- [x] `system` — `system.py` → `SystemHandler` + `@quantity("system") System` | test: `test_system.py` ✅ +- [x] `run_info` — `run_info.py` → `RunInfoHandler` + `@quantity("run_info") RunInfo` | test: `test_run_info.py` ✅ +- [x] `polarization` — `polarization.py` → `PolarizationHandler` + `@quantity("polarization") Polarization` | test: `test_polarization.py` ✅ +- [x] `piezoelectric_tensor` — `piezoelectric_tensor.py` → `PiezoelectricTensorHandler` + `@quantity("piezoelectric_tensor") PiezoelectricTensor` | test: `test_piezoelectric_tensor.py` ✅ +- [x] `dielectric_tensor` — `dielectric_tensor.py` → `DielectricTensorHandler` + `@quantity("dielectric_tensor") DielectricTensor` | test: `test_dielectric_tensor.py` ✅ +- [x] `elastic_modulus` — `elastic_modulus.py` → `ElasticModulusHandler` + `@quantity("elastic_modulus") ElasticModulus` | test: `test_elastic_modulus.py` ✅ +- [x] `workfunction` — `workfunction.py` → `WorkfunctionHandler` + `@quantity("workfunction") Workfunction` | test: `test_workfunction.py` ✅ +- [x] `dielectric_function` — `dielectric_function.py` → `DielectricFunctionHandler` + `@quantity("dielectric_function") DielectricFunction` | test: `test_dielectric_function.py` ✅ +- [x] `effective_coulomb` — `effective_coulomb.py` → `EffectiveCoulombHandler` + `@quantity("effective_coulomb") EffectiveCoulomb` | test: `test_effective_coulomb.py` ✅ +- [x] `internal_strain` — `internal_strain.py` → `InternalStrainHandler` + `@quantity("internal_strain") InternalStrain` | test: `test_internal_strain.py` ✅ + - Note: uses `structure.Mixin`; call `StructureHandler.from_data` directly in the Handler. +- [x] `born_effective_charge` — `born_effective_charge.py` → `BornEffectiveChargeHandler` + `@quantity("born_effective_charge") BornEffectiveCharge` | test: `test_born_effective_charge.py` ✅ + - Note: uses `structure.Mixin`; requires `structure` to be ported first (or compose directly via `StructureHandler`). +- [x] `force_constant` — `force_constant.py` → `ForceConstantHandler` + `@quantity("force_constant") ForceConstant` | test: `test_force_constant.py` ✅ + - Note: uses `structure.Mixin`; same composition pattern as above. + +--- + +### Group 2 — Step-indexed (no external composition) + +These use `slice_.Mixin` and require `slice_steps` in the Handler. + +- [x] `energy` — `energy.py` → `EnergyHandler` + `@quantity("energy") Energy(graph.Mixin)` | test: `test_energy.py` ✅ +- [x] `electronic_minimization` — `electronic_minimization.py` → `ElectronicMinimizationHandler` + `@quantity("electronic_minimization") ElectronicMinimization(graph.Mixin)` | test: `test_electronic_minimization.py` ✅ +- [x] `pair_correlation` — `pair_correlation.py` → `PairCorrelationHandler` + `@quantity("pair_correlation") PairCorrelation(graph.Mixin)` | test: `test_pair_correlation.py` ✅ + +--- + +### Group 3 — Internal helpers (port before their dependents) + +These are not in `QUANTITIES`/`GROUPS` but are composed into other quantities. + +- [x] `_stoichiometry` — `_stoichiometry.py` → `StoichiometryHandler` added | test: `test_stoichiometry.py` + - Note: internal, no `@quantity` registration needed; stays in `QUANTITIES` as `"_stoichiometry"` until all dependents are ported. +- [x] `cell` — `cell.py` → `CellHandler` added | no dedicated public test (internal helper) + - Note: not in `QUANTITIES`; used by `StructureHandler` via composition. + +--- + +### Group 4 — Structure (central dependency) + +Port before any quantity that composes Structure. + +- [x] `structure` — `structure.py` → `StructureHandler` added | test: `test_structure.py` + - Note: step-indexed; composes `CellHandler` and `StoichiometryHandler` internally. + - Note: many downstream quantities depend on `StructureHandler.from_data`. + +--- + +### Group 5 — Step-indexed + structure composition + +Requires `structure` to be ported first. + +- [x] `force` — `force.py` → `ForceHandler` + `@quantity("force") Force(view.Mixin)` | test: `test_force.py` ✅ +- [x] `stress` — `stress.py` → `StressHandler` + `@quantity("stress") Stress` | test: `test_stress.py` ✅ +- [x] `velocity` — `velocity.py` → `VelocityHandler` + `@quantity("velocity") Velocity(view.Mixin)` | test: `test_velocity.py` ✅ +- [x] `local_moment` — `local_moment.py` → `LocalMomentHandler` + `@quantity("local_moment") LocalMoment(view.Mixin)` | test: `test_local_moment.py` ✅ +- [x] `nics` — `nics.py` → `NicsHandler` + `@quantity("nics") Nics(view.Mixin)` | test: `test_nics.py` ✅ +- [x] `density` — `density.py` → `DensityHandler` + `@quantity("density") Density(view.Mixin)` | test: `test_density.py` ✅ +- [x] `partial_density` — `partial_density.py` → `PartialDensityHandler` + `@quantity("partial_density") PartialDensity(view.Mixin)` | test: `test_partial_density.py` ✅ +- [x] `potential` — `potential.py` → `PotentialHandler` + `@quantity("potential") Potential(view.Mixin)` | test: `test_potential.py` ✅ +- [x] `current_density` — `current_density.py` → `CurrentDensityHandler` + `@quantity("current_density") CurrentDensity` | test: `test_current_density.py` ✅ +- [x] `_CONTCAR` — `_CONTCAR.py` → `CONTCARHandler` + `@quantity("_CONTCAR") CONTCAR(view.Mixin)` | test: `test_contcar.py` ✅ + +--- + +### Group 6 — Projector / kpoint / dispersion chain + +Port in order: `kpoint` → `_dispersion` → `projector` → `dos` → `band`. + +- [x] `kpoint` — `kpoint.py` → `KpointHandler` + `@quantity("kpoint") Kpoint` | test: `test_kpoint.py` ✅ +- [x] `_dispersion` — `_dispersion.py` → `DispersionHandler` + `@quantity("_dispersion") Dispersion` | test: `test_dispersion.py` ✅ +- [x] `projector` — `projector.py` → `ProjectorHandler` + `@quantity("projector") Projector` | test: `test_projector.py` ✅ +- [x] `dos` — `dos.py` → `DosHandler` + `@quantity("dos") Dos(graph.Mixin)` | test: `test_dos.py` ✅ +- [x] `band` — `band.py` → `BandHandler` + `@quantity("band") Band(graph.Mixin)` | test: `test_band.py` ✅ + - Note: most complex standalone quantity; composes `ProjectorHandler`, `KpointHandler`, `DispersionHandler`. + +--- + +### Group 7 — Phonon group + +Port `phonon.py` (the shared Mixin) conceptually before the three group members. + +- [x] `phonon.mode` — `phonon_mode.py` → `PhononModeHandler` + `@quantity("mode", group="phonon") PhononMode` | test: `test_phonon_mode.py` ✅ +- [x] `phonon.band` — `phonon_band.py` → `PhononBandHandler` + `@quantity("band", group="phonon") PhononBand(graph.Mixin)` | test: `test_phonon_band.py` ✅ +- [x] `phonon.dos` — `phonon_dos.py` → `PhononDosHandler` + `@quantity("dos", group="phonon") PhononDos(graph.Mixin)` | test: `test_phonon_dos.py` ✅ + +--- + +### Group 8 — Exciton group + +- [x] `exciton.eigenvector` — `exciton_eigenvector.py` → `ExcitonEigenvectorHandler` + `@quantity("eigenvector", group="exciton") ExcitonEigenvector` | test: `test_exciton_eigenvector.py` ✅ +- [x] `exciton.density` — `exciton_density.py` → `ExcitonDensityHandler` + `@quantity("density", group="exciton") ExcitonDensity(view.Mixin)` | test: `test_exciton_density.py` ✅ + +--- + +### Group 9 — Electron-phonon group (complex) ✅ + +`ElectronPhononSelfEnergy` and `ElectronPhononTransport` also inherit `abc.Sequence` — figure out how to expose `__len__`/`__getitem__` on the Dispatcher without using the Refinery pattern. + +- [x] `electron_phonon.chemical_potential` — `electron_phonon_chemical_potential.py` → `ElectronPhononChemicalPotential(Refinery)` | test: `test_electron_phonon_chemical_potential.py` + - Group: `@quantity("chemical_potential", group="electron_phonon")`; remove from `GROUPS["electron_phonon"]`. +- [x] `electron_phonon.bandgap` — `electron_phonon_bandgap.py` → `ElectronPhononBandgap(Refinery, abc.Sequence)` | test: `test_electron_phonon_bandgap.py` + - Note: step-indexed via `abc.Sequence`; complex internal structure. + - Group: `@quantity("bandgap", group="electron_phonon")`. +- [x] `electron_phonon.self_energy` — `electron_phonon_self_energy.py` → `ElectronPhononSelfEnergy(Refinery, abc.Sequence)` | test: `test_electron_phonon_self_energy.py` + - Note: `abc.Sequence` exposes sub-instances; requires careful Handler factory design. + - Group: `@quantity("self_energy", group="electron_phonon")`. +- [x] `electron_phonon.transport` — `electron_phonon_transport.py` → `ElectronPhononTransport(Refinery, abc.Sequence, graph.Mixin)` | test: `test_electron_phonon_transport.py` + - Note: most complex quantity; `abc.Sequence` + graph; leave for last. + - Group: `@quantity("transport", group="electron_phonon")`. + +--- + +## Remaining Work (post-porting) + +All 40+ quantities have been ported to Dispatcher/Handler. These phases complete the migration. + +--- + +### Phase 1: Add `from_path`/`from_file` to all Dispatchers (unblock `test_factory_methods`) ✅ + +**Problem**: ~20 dispatcher classes lack `from_path`/`from_file`. The `check_factory_methods` test fixture calls `cls.from_path()` directly, so all 38 tests are skipped with `"Dispatcher not yet wired to Calculation"`. + +**Approach**: Modify the `@quantity` decorator in `dispatch.py` to auto-inject `from_path`/`from_file` classmethods. All dispatchers accept `(source, quantity_name)` — just need `FileSource` construction. Remove any manually-defined `from_path`/`from_file` from dispatcher classes to avoid duplication. + +- [x] Remove existing `from_path`/`from_file` from all dispatcher classes that define them manually +- [x] Inject `from_path`/`from_file` via `@quantity` decorator in `dispatch.py` +- [x] Remove `@pytest.mark.skip(reason="Dispatcher not yet wired to Calculation")` from all 38 test files +- [x] Fix `__str__` on all dispatchers to accept `selection=None` and forward to `merge_strings` +- [x] Fix `_dispatch()` to handle `remaining_selection=None` with default-valued handler params +- [x] Fix conftest: derive quantity from `_quantity_name`, handle methods without `selection`, relax assertions +- [x] `pytest tests/calculation/ -k test_factory_methods` — 42 passed ✅ +- [x] `pytest tests/calculation/` — 2051 passed, 6 failed (all mdtraj, pre-existing) ✅ + +--- + +### Phase 2: Port `Calculation.to_database()` to new architecture + +**Problem**: `_compute_database_data()` iterates the old `QUANTITIES` tuple (only "structure", "_stoichiometry") and `GROUPS` dict (empty). It calls `Refinery._read_to_database()` which new dispatchers don't have. + +- [ ] Add `_read_to_database()` on each Dispatcher (delegates to Handler.to_database() via `_dispatch()`) +- [ ] Rewrite `_compute_database_data()` to iterate `_REGISTRY` instead of `QUANTITIES`/`GROUPS` +- [ ] Port selection-handling and caching logic from `Refinery._read_to_database()` into new system +- [ ] Handle grouped quantities (phonon.band, electron_phonon.self_energy, etc.) via Group wrapper +- [ ] Existing database tests pass; `Calculation.to_database()` produces same output + +--- + +### Phase 3: Remove old Refinery classes + +*Depends on Phase 2.* + +- [ ] Delete `class Stoichiometry(base.Refinery)` from `_stoichiometry.py` (Handler already exists) +- [ ] Delete `class Cell(slice_.Mixin, base.Refinery)` from `cell.py` (Handler already exists) +- [ ] Delete `class Structure(slice_.Mixin, base.Refinery, view.Mixin)` from `structure.py` (Handler already exists) +- [ ] Remove `QUANTITIES = ("structure", "_stoichiometry")` from `__init__.py` +- [ ] Remove `_add_all_refinement_classes(Calculation)` call and helpers (`_make_property`, `_make_group`) +- [ ] `pytest tests/ -x` passes + +--- + +### Phase 4: Remove dead infrastructure + +*Depends on Phase 3.* + +- [ ] Delete `src/py4vasp/_calculation/base.py` (Refinery, `data_access` decorator) +- [ ] Remove `phonon.Mixin` from `phonon.py` (still uses `@base.data_access`) +- [ ] Remove `slice_.Mixin` class (keep `slice_.examples` decorator for docstrings) +- [ ] Delete `src/py4vasp/_calculation/local_moment.py~` (backup file) +- [ ] Clean up dead imports across all files +- [ ] Update `__all__` in `__init__.py` to expose dispatcher names from `_REGISTRY` + +--- + +### Phase 5: Final verification + +- [ ] `pytest tests/ -x` — full suite green +- [ ] `grep -r "Refinery" src/` — zero hits +- [ ] `grep -r "data_access" src/` — zero hits +- [ ] `grep -r "skip.*factory" tests/` — zero hits +- [ ] `Calculation.to_database()` works end-to-end diff --git a/pytest_output.txt b/pytest_output.txt new file mode 100644 index 000000000..6c7ed8856 --- /dev/null +++ b/pytest_output.txt @@ -0,0 +1,410 @@ +============================= test session starts ============================= +platform win32 -- Python 3.12.7, pytest-9.0.1, pluggy-1.6.0 -- C:\Users\mschl\miniconda3\envs\p4v\python.exe +cachedir: .pytest_cache +hypothesis profile 'default' +rootdir: C:\Users\mschl\Documents\code\python\py4vasp +configfile: pyproject.toml +plugins: anyio-4.12.0, hypothesis-6.148.5, cov-7.0.0, timeout-2.4.0 +collecting ... collected 204 items + +tests/calculation/test_electron_phonon_chemical_potential.py::test_read[carrier_den] PASSED [ 0%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_read[mu] PASSED [ 0%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_read[carrier_per_cell] PASSED [ 1%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_label[carrier_den] PASSED [ 1%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_label[mu] PASSED [ 2%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_label[carrier_per_cell] PASSED [ 2%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_mu_tag[carrier_den] PASSED [ 3%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_mu_tag[mu] PASSED [ 3%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_mu_tag[carrier_per_cell] PASSED [ 4%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_print[carrier_den] PASSED [ 4%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_print[mu] PASSED [ 5%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_print[carrier_per_cell] PASSED [ 5%] +tests/calculation/test_electron_phonon_chemical_potential.py::test_factory_methods SKIPPED [ 6%] +tests/calculation/test_electron_phonon_bandgap.py::test_len[nonpolarized] PASSED [ 6%] +tests/calculation/test_electron_phonon_bandgap.py::test_len[collinear] PASSED [ 7%] +tests/calculation/test_electron_phonon_bandgap.py::test_indexing_and_iteration[nonpolarized] FAILED [ 7%] +tests/calculation/test_electron_phonon_bandgap.py::test_indexing_and_iteration[collinear] FAILED [ 8%] +tests/calculation/test_electron_phonon_bandgap.py::test_read_mapping[nonpolarized] PASSED [ 8%] +tests/calculation/test_electron_phonon_bandgap.py::test_read_mapping[collinear] PASSED [ 9%] +tests/calculation/test_electron_phonon_bandgap.py::test_read_instance[nonpolarized] PASSED [ 9%] +tests/calculation/test_electron_phonon_bandgap.py::test_read_instance[collinear] PASSED [ 10%] +tests/calculation/test_electron_phonon_bandgap.py::test_read_instance_metadata[nonpolarized] PASSED [ 10%] +tests/calculation/test_electron_phonon_bandgap.py::test_read_instance_metadata[collinear] PASSED [ 11%] +tests/calculation/test_electron_phonon_bandgap.py::test_plot_instance[nonpolarized] PASSED [ 11%] +tests/calculation/test_electron_phonon_bandgap.py::test_plot_instance[collinear] PASSED [ 12%] +tests/calculation/test_electron_phonon_bandgap.py::test_plot_multiple_selections[nonpolarized] PASSED [ 12%] +tests/calculation/test_electron_phonon_bandgap.py::test_plot_multiple_selections[collinear] PASSED [ 13%] +tests/calculation/test_electron_phonon_bandgap.py::test_plot_direct_gap_renormalization[nonpolarized] PASSED [ 13%] +tests/calculation/test_electron_phonon_bandgap.py::test_plot_direct_gap_renormalization[collinear] PASSED [ 14%] +tests/calculation/test_electron_phonon_bandgap.py::test_selections[nonpolarized-carrier_den] PASSED [ 14%] +tests/calculation/test_electron_phonon_bandgap.py::test_selections[nonpolarized-carrier_per_cell] PASSED [ 15%] +tests/calculation/test_electron_phonon_bandgap.py::test_selections[nonpolarized-mu] PASSED [ 15%] +tests/calculation/test_electron_phonon_bandgap.py::test_selections[collinear-carrier_den] PASSED [ 16%] +tests/calculation/test_electron_phonon_bandgap.py::test_selections[collinear-carrier_per_cell] PASSED [ 16%] +tests/calculation/test_electron_phonon_bandgap.py::test_selections[collinear-mu] PASSED [ 17%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_returns_instances[nonpolarized-nbands_sum] PASSED [ 17%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_returns_instances[nonpolarized-selfen_delta] PASSED [ 18%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_returns_instances[nonpolarized-selfen_carrier_den] PASSED [ 18%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_returns_instances[collinear-nbands_sum] PASSED [ 19%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_returns_instances[collinear-selfen_delta] PASSED [ 19%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_returns_instances[collinear-selfen_carrier_den] PASSED [ 20%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_non_SERTA_returns_empty[nonpolarized] PASSED [ 20%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_non_SERTA_returns_empty[collinear] PASSED [ 21%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_multiple[nonpolarized] PASSED [ 21%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_multiple[collinear] PASSED [ 22%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_nested[nonpolarized] PASSED [ 22%] +tests/calculation/test_electron_phonon_bandgap.py::test_select_nested[collinear] PASSED [ 23%] +tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[nonpolarized-invalid_selection=0.01] FAILED [ 23%] +tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[nonpolarized-nbands_sum:0.01] PASSED [ 24%] +tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[nonpolarized-selfen_delta] PASSED [ 24%] +tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[collinear-invalid_selection=0.01] FAILED [ 25%] +tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[collinear-nbands_sum:0.01] PASSED [ 25%] +tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[collinear-selfen_delta] PASSED [ 25%] +tests/calculation/test_electron_phonon_bandgap.py::test_print_mapping[nonpolarized] FAILED [ 26%] +tests/calculation/test_electron_phonon_bandgap.py::test_print_mapping[collinear] FAILED [ 26%] +tests/calculation/test_electron_phonon_bandgap.py::test_print_instance[nonpolarized] PASSED [ 27%] +tests/calculation/test_electron_phonon_bandgap.py::test_print_instance[collinear] PASSED [ 27%] +tests/calculation/test_electron_phonon_bandgap.py::test_factory_methods SKIPPED [ 28%] +tests/calculation/test_electron_phonon_self_energy.py::test_len FAILED [ 28%] +tests/calculation/test_electron_phonon_self_energy.py::test_indexing_and_iteration FAILED [ 29%] +tests/calculation/test_electron_phonon_self_energy.py::test_read_mapping PASSED [ 29%] +tests/calculation/test_electron_phonon_self_energy.py::test_read_instance PASSED [ 30%] +tests/calculation/test_electron_phonon_self_energy.py::test_read_instance_metadata PASSED [ 30%] +tests/calculation/test_electron_phonon_self_energy.py::test_selections[carrier_den] PASSED [ 31%] +tests/calculation/test_electron_phonon_self_energy.py::test_selections[carrier_per_cell] PASSED [ 31%] +tests/calculation/test_electron_phonon_self_energy.py::test_selections[mu] PASSED [ 32%] +tests/calculation/test_electron_phonon_self_energy.py::test_select_returns_instances[nbands_sum] PASSED [ 32%] +tests/calculation/test_electron_phonon_self_energy.py::test_select_returns_instances[selfen_delta] PASSED [ 33%] +tests/calculation/test_electron_phonon_self_energy.py::test_select_returns_instances[selfen_carrier_den] PASSED [ 33%] +tests/calculation/test_electron_phonon_self_energy.py::test_select_returns_instances[scattering_approx] PASSED [ 34%] +tests/calculation/test_electron_phonon_self_energy.py::test_select_multiple PASSED [ 34%] +tests/calculation/test_electron_phonon_self_energy.py::test_select_nested PASSED [ 35%] +tests/calculation/test_electron_phonon_self_energy.py::test_incorrect_selection[invalid_selection=0.01] FAILED [ 35%] +tests/calculation/test_electron_phonon_self_energy.py::test_incorrect_selection[nbands_sum:0.01] PASSED [ 36%] +tests/calculation/test_electron_phonon_self_energy.py::test_incorrect_selection[selfen_delta] PASSED [ 36%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_valid_and_invalid_access[valid_indices0-10.0] PASSED [ 37%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_valid_and_invalid_access[valid_indices1-20.0] PASSED [ 37%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_valid_and_invalid_access[valid_indices2-30.0] PASSED [ 38%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_negative_indices[negative_indices0-30.0] PASSED [ 38%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_negative_indices[negative_indices1-10.0] PASSED [ 39%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_negative_indices[negative_indices2-20.0] PASSED [ 39%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_valid_bands PASSED [ 40%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_invalid_access[invalid_indices0] PASSED [ 40%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_invalid_access[invalid_indices1] PASSED [ 41%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_invalid_access[invalid_indices2] PASSED [ 41%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_out_of_bounds_access[out_of_bounds_indices0] PASSED [ 42%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_out_of_bounds_access[out_of_bounds_indices1] PASSED [ 42%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_incorrect_arguments PASSED [ 43%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_should_only_succeed_if_index_is_valid PASSED [ 43%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_self_energy[fan] PASSED [ 44%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_self_energy[debye_waller] PASSED [ 44%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_self_energy[self_energy] PASSED [ 45%] +tests/calculation/test_electron_phonon_self_energy.py::test_sparse_tensor_self_energy[energies] PASSED [ 45%] +tests/calculation/test_electron_phonon_self_energy.py::test_eigenvalues PASSED [ 46%] +tests/calculation/test_electron_phonon_self_energy.py::test_print_mapping FAILED [ 46%] +tests/calculation/test_electron_phonon_self_energy.py::test_print_instance PASSED [ 47%] +tests/calculation/test_electron_phonon_self_energy.py::test_factory_methods SKIPPED [ 47%] +tests/calculation/test_electron_phonon_transport.py::test_len FAILED [ 48%] +tests/calculation/test_electron_phonon_transport.py::test_indexing_and_iteration FAILED [ 48%] +tests/calculation/test_electron_phonon_transport.py::test_read_mapping PASSED [ 49%] +tests/calculation/test_electron_phonon_transport.py::test_read_instance PASSED [ 49%] +tests/calculation/test_electron_phonon_transport.py::test_read_instance_metadata PASSED [ 50%] +tests/calculation/test_electron_phonon_transport.py::test_read_instance_metadata_CRTA PASSED [ 50%] +tests/calculation/test_electron_phonon_transport.py::test_selections[carrier_den] PASSED [ 50%] +tests/calculation/test_electron_phonon_transport.py::test_selections[carrier_per_cell] PASSED [ 51%] +tests/calculation/test_electron_phonon_transport.py::test_selections[mu] PASSED [ 51%] +tests/calculation/test_electron_phonon_transport.py::test_selections_CRTA PASSED [ 52%] +tests/calculation/test_electron_phonon_transport.py::test_select_returns_instances[nbands_sum] PASSED [ 52%] +tests/calculation/test_electron_phonon_transport.py::test_select_returns_instances[selfen_delta] PASSED [ 53%] +tests/calculation/test_electron_phonon_transport.py::test_select_returns_instances[selfen_carrier_den] PASSED [ 53%] +tests/calculation/test_electron_phonon_transport.py::test_select_returns_instances[scattering_approx] PASSED [ 54%] +tests/calculation/test_electron_phonon_transport.py::test_select_multiple PASSED [ 54%] +tests/calculation/test_electron_phonon_transport.py::test_select_nested PASSED [ 55%] +tests/calculation/test_electron_phonon_transport.py::test_select_missing PASSED [ 55%] +tests/calculation/test_electron_phonon_transport.py::test_incorrect_selection[invalid_selection=0.01] FAILED [ 56%] +tests/calculation/test_electron_phonon_transport.py::test_incorrect_selection[nbands_sum:0.01] PASSED [ 56%] +tests/calculation/test_electron_phonon_transport.py::test_incorrect_selection[selfen_delta] PASSED [ 57%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping[electronic_conductivity] PASSED [ 57%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping[mobility] PASSED [ 58%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping[seebeck] PASSED [ 58%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping[peltier] PASSED [ 59%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping[electronic_thermal_conductivity] PASSED [ 59%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[None-expected0] PASSED [ 60%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[isotropic-expected1] PASSED [ 60%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[xx-0] PASSED [ 61%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[yy-4] PASSED [ 61%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[zz-8] PASSED [ 62%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[xy-expected5] PASSED [ 62%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[xz-expected6] PASSED [ 63%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_with_direction[yz-expected7] PASSED [ 63%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_multiple_directions PASSED [ 64%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_select_scattering_approx PASSED [ 64%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_select_temperature[T=300, T=400] PASSED [ 65%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_select_temperature[temperature=300, temperature=400] PASSED [ 65%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_complex_selection[mobility(zz(nbands_sum=32(T=200)))] PASSED [ 66%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_complex_selection[T=200(mobility(zz(nbands_sum=32)))] PASSED [ 66%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_complex_selection[zz(T=200(mobility(nbands_sum=32)))] PASSED [ 67%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_complex_selection[nbands_sum=32(zz(T=200(mobility)))] PASSED [ 67%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_chemical_potential[carrier_den] PASSED [ 68%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_chemical_potential[carrier_per_cell] PASSED [ 68%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_chemical_potential[mu] PASSED [ 69%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_CRTA PASSED [ 69%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_incorrect_selection[unknown_selection] PASSED [ 70%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_incorrect_selection[mobility(seebeck)] PASSED [ 70%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_incorrect_selection[seebeck, peltier] PASSED [ 71%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_incorrect_selection[mobility(xx(yy))] PASSED [ 71%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance[electronic_conductivity] PASSED [ 72%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance[mobility] PASSED [ 72%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance[seebeck] PASSED [ 73%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance[peltier] PASSED [ 73%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance[electronic_thermal_conductivity] PASSED [ 74%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[None-expected0] PASSED [ 74%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[isotropic-expected1] PASSED [ 75%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[xx-0] PASSED [ 75%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[yy-4] PASSED [ 75%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[zz-8] PASSED [ 76%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[xy-expected5] PASSED [ 76%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[xz-expected6] PASSED [ 77%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_direction[yz-expected7] PASSED [ 77%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_multiple_directions PASSED [ 78%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_incorrect_selection[unknown_selection] PASSED [ 78%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_incorrect_selection[seebeck(peltier)] PASSED [ 79%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_incorrect_selection[seebeck, peltier] PASSED [ 79%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_incorrect_selection[seebeck(xx(yy))] PASSED [ 80%] +tests/calculation/test_electron_phonon_transport.py::test_figure_of_merit PASSED [ 80%] +tests/calculation/test_electron_phonon_transport.py::test_figure_of_merit_with_argument[1.0] PASSED [ 81%] +tests/calculation/test_electron_phonon_transport.py::test_figure_of_merit_with_argument[kappa_lattice1] PASSED [ 81%] +tests/calculation/test_electron_phonon_transport.py::test_figure_of_merit_with_wrong_size_argument PASSED [ 82%] +tests/calculation/test_electron_phonon_transport.py::test_print_mapping FAILED [ 82%] +tests/calculation/test_electron_phonon_transport.py::test_print_instance PASSED [ 83%] +tests/calculation/test_electron_phonon_transport.py::test_factory_methods SKIPPED [ 83%] +tests/calculation/test_electron_phonon_transport.py::test_read_instance_spin PASSED [ 84%] +tests/calculation/test_electron_phonon_transport.py::test_read_instance_source_selection_forwarded PASSED [ 84%] +tests/calculation/test_electron_phonon_transport.py::test_selections_spin PASSED [ 85%] +tests/calculation/test_electron_phonon_transport.py::test_selections_no_spin PASSED [ 85%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_default_sums[electronic_conductivity] PASSED [ 86%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_default_sums[mobility] PASSED [ 86%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_default_sums[seebeck] PASSED [ 87%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_default_sums[peltier] PASSED [ 87%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_default_sums[electronic_thermal_conductivity] PASSED [ 88%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_selection[up-0] PASSED [ 88%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_selection[down-1] PASSED [ 89%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_and_direction PASSED [ 89%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_both_spins PASSED [ 90%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_default_sums[electronic_conductivity] PASSED [ 90%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_default_sums[mobility] PASSED [ 91%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_default_sums[seebeck] PASSED [ 91%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_default_sums[peltier] PASSED [ 92%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_default_sums[electronic_thermal_conductivity] PASSED [ 92%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_selection[up-0] PASSED [ 93%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_selection[down-1] PASSED [ 93%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_source_prefix_with_spin_data[electronic_conductivity] PASSED [ 94%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_source_prefix_with_spin_data[mobility] PASSED [ 94%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_source_prefix_with_spin_data[seebeck] PASSED [ 95%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_source_prefix_with_spin_data[peltier] PASSED [ 95%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_source_prefix_with_spin_data[electronic_thermal_conductivity] PASSED [ 96%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_source_prefix_non_spin_data_raises PASSED [ 96%] +tests/calculation/test_electron_phonon_transport.py::test_plot_no_spin_data_spin_selection_raises PASSED [ 97%] +tests/calculation/test_electron_phonon_transport.py::test_plot_no_spin_data_mapping_spin_selection_raises PASSED [ 97%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_mobility_selection[up-0] PASSED [ 98%] +tests/calculation/test_electron_phonon_transport.py::test_plot_instance_spin_mobility_selection[down-1] PASSED [ 98%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_mobility_selection[up-0] PASSED [ 99%] +tests/calculation/test_electron_phonon_transport.py::test_plot_mapping_spin_mobility_selection[down-1] PASSED [ 99%] +tests/calculation/test_electron_phonon_transport.py::test_plot_no_spin_data_mobility_spin_selection_raises PASSED [100%] + +================================== FAILURES =================================== +__________________ test_indexing_and_iteration[nonpolarized] __________________ +tests\calculation\test_electron_phonon_bandgap.py:155: in test_indexing_and_iteration + assert instance.parent is band_gap +E assert is +E + where = .parent +___________________ test_indexing_and_iteration[collinear] ____________________ +tests\calculation\test_electron_phonon_bandgap.py:155: in test_indexing_and_iteration + assert instance.parent is band_gap +E assert is +E + where = .parent +________ test_incorrect_selection[nonpolarized-invalid_selection=0.01] ________ +tests\calculation\test_electron_phonon_bandgap.py:290: in test_incorrect_selection + band_gap.select(selection) +src\py4vasp\_calculation\electron_phonon_bandgap.py:281: in select + return self._handler_factory(raw).select(selection) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_bandgap.py:161: in select + indices = self._accumulator().select_indices( +src\py4vasp\_calculation\electron_phonon_accumulator.py:70: in select_indices + for index_ in self._filter_indices(selection, kwargs_filters) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:82: in _filter_indices + remaining_indices = list(remaining_indices) + ^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:94: in _filter_assignment + if self._match_key_value(index_, key, str(value)): + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:98: in _match_key_value + instance_value = self.get_data(key, index_) + ^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:114: in get_data + self._raise_error_if_not_present(name, expected_name=mu_tag) +src\py4vasp\_calculation\electron_phonon_accumulator.py:122: in _raise_error_if_not_present + valid_names.remove(self._name) +E KeyError: 'electron_phonon_bandgap_handler' +_________ test_incorrect_selection[collinear-invalid_selection=0.01] __________ +tests\calculation\test_electron_phonon_bandgap.py:290: in test_incorrect_selection + band_gap.select(selection) +src\py4vasp\_calculation\electron_phonon_bandgap.py:281: in select + return self._handler_factory(raw).select(selection) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_bandgap.py:161: in select + indices = self._accumulator().select_indices( +src\py4vasp\_calculation\electron_phonon_accumulator.py:70: in select_indices + for index_ in self._filter_indices(selection, kwargs_filters) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:82: in _filter_indices + remaining_indices = list(remaining_indices) + ^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:94: in _filter_assignment + if self._match_key_value(index_, key, str(value)): + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:98: in _match_key_value + instance_value = self.get_data(key, index_) + ^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:114: in get_data + self._raise_error_if_not_present(name, expected_name=mu_tag) +src\py4vasp\_calculation\electron_phonon_accumulator.py:122: in _raise_error_if_not_present + valid_names.remove(self._name) +E KeyError: 'electron_phonon_bandgap_handler' +______________________ test_print_mapping[nonpolarized] _______________________ +tests\calculation\test_electron_phonon_bandgap.py:295: in test_print_mapping + assert re.search(band_gap.ref.mapping_pattern, str(band_gap), re.MULTILINE) +E assert None +E + where None = ('Electron-phonon bandgap with 4 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]', "Electron-phonon bandgap handler with 4 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-10.32472187 -2.04940455 -1.76249068 2.18068674 2.66918125]\n scattering_approx: ['MRTA_TAU' 'SERTA']", re.MULTILINE) +E + where = re.search +E + and 'Electron-phonon bandgap with 4 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]' = namespace(naccumulators=4, indices=array([0, 1, 3, 4]), fundamental=VaspData(array([[ -0.69872384],\n [ 15.58191799],\n [ 18.25777394],\n [ -8.13746226],\n [-23.00295462]])), fundamental_renorm=VaspData(array([[[-11.6337025 , -4.75260521, 12.07693774, 5.84548688,\n 1.08609908, -2.60714231]],\n\n [[-15.03548617, -11.31805403, -13.29097232, 3.33615862,\n -20.08308288, -0.51522233]],\n\n [[ 0.6111418 , 0.14682233, 6.97069027, 4.9531376 ,\n 8.74394084, -10.71490414]],\n\n [[-10.11716136, -5.45067988, 15.25316052, 17.3255197 ,\n -17.69116922, 17.14170855]],\n\n [[ 19.16866672, -1.94130858, -7.67335238, -0.98434872,\n 8.16793369, -1.32801036]]])), direct=VaspData(array([[ 2.50305504],\n [-6.98645358],\n [-5.53612536],\n [ 1.25674325],\n [-2.71784768]])), direct_renorm=VaspData(array([[[-18.09883575, -4.05034335, -6.73639966, -5.22773509,\n 24.23612535, 4.4686224 ]],\n\n [[-17.21288603, 7.98867433, 19.12149503, -1.3036616 ,\n -9.06741626, -6.57431361]],\n\n [[ 13.44805744, 13.72748006, 4.4024852 , -15.78610124,\n 13.0366001 ...} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n\nFundamental gap:\n Temperature \\(K\\) KS gap \\(eV\\) QP gap \\(eV\\) KS-QP gap \\(meV\\)\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}').mapping_pattern +E + where namespace(naccumulators=4, indices=array([0, 1, 3, 4]), fundamental=VaspData(array([[ -0.69872384],\n [ 15.58191799],\n [ 18.25777394],\n [ -8.13746226],\n [-23.00295462]])), fundamental_renorm=VaspData(array([[[-11.6337025 , -4.75260521, 12.07693774, 5.84548688,\n 1.08609908, -2.60714231]],\n\n [[-15.03548617, -11.31805403, -13.29097232, 3.33615862,\n -20.08308288, -0.51522233]],\n\n [[ 0.6111418 , 0.14682233, 6.97069027, 4.9531376 ,\n 8.74394084, -10.71490414]],\n\n [[-10.11716136, -5.45067988, 15.25316052, 17.3255197 ,\n -17.69116922, 17.14170855]],\n\n [[ 19.16866672, -1.94130858, -7.67335238, -0.98434872,\n 8.16793369, -1.32801036]]])), direct=VaspData(array([[ 2.50305504],\n [-6.98645358],\n [-5.53612536],\n [ 1.25674325],\n [-2.71784768]])), direct_renorm=VaspData(array([[[-18.09883575, -4.05034335, -6.73639966, -5.22773509,\n 24.23612535, 4.4686224 ]],\n\n [[-17.21288603, 7.98867433, 19.12149503, -1.3036616 ,\n -9.06741626, -6.57431361]],\n\n [[ 13.44805744, 13.72748006, 4.4024852 , -15.78610124,\n 13.0366001 ...} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n\nFundamental gap:\n Temperature \\(K\\) KS gap \\(eV\\) QP gap \\(eV\\) KS-QP gap \\(meV\\)\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}') = .ref +E + and "Electron-phonon bandgap handler with 4 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-10.32472187 -2.04940455 -1.76249068 2.18068674 2.66918125]\n scattering_approx: ['MRTA_TAU' 'SERTA']" = str() +E + and re.MULTILINE = re.MULTILINE +________________________ test_print_mapping[collinear] ________________________ +tests\calculation\test_electron_phonon_bandgap.py:295: in test_print_mapping + assert re.search(band_gap.ref.mapping_pattern, str(band_gap), re.MULTILINE) +E assert None +E + where None = ('Electron-phonon bandgap with 4 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]', "Electron-phonon bandgap handler with 4 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-10.32472187 -2.04940455 -1.76249068 2.18068674 2.66918125]\n scattering_approx: ['MRTA_TAU' 'SERTA']", re.MULTILINE) +E + where = re.search +E + and 'Electron-phonon bandgap with 4 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]' = namespace(naccumulators=4, indices=array([0, 1, 3, 4]), fundamental=VaspData(array([[ -3.44616154, -19.7728207 , 22.11833881],\n [ -8.54658364, -18.30250658, -15.54715963],\n [ -0.42216143, 8.6986917 , -7.30986158],\n [ -6.76019621, 1.50705188, -1.2590991 ],\n [ 1.09588539, 0.62770383, 2.89870303]])), fundamental_renorm=VaspData(array([[[ 1.31447368, 2.82453526, -2.94655269, 0.31682863,\n 19.32203649, 2.29440329],\n [ 0.93448729, 2.57937955, 1.19506417, 1.30644151,\n 8.56364231, -0.27896706],\n [ 1.03527326, 6.10915995, 9.01056266, -8.2157778 ,\n 5.22714501, 3.18536412]],\n\n [[ 5.48283599, 4.19715567, 4.44685687, 6.43824551,\n -5.52373766, 8.18635823],\n [-13.55073087, 1.22158398, 13.60890006, -16.40893432,\n -26.9191767 , 3.2260027 ],\n [ 28.24687453, 15.48433366, 2.05068837, -15.5690078 ,\n 3.02198295, -4.42156019]],\n\n [[ 0.46848899, 9.64431376, -10.9034597 , 7.77889988,\n -10.92221812, 17.93433268],\n [ -5.97143202, 5.60993696, 8.29856501, -2.96932009,\n 6.38247299, 6.95287228],\n ...} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n\nFundamental gap:\n Temperature \\(K\\) KS gap \\(eV\\) QP gap \\(eV\\) KS-QP gap \\(meV\\)\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}').mapping_pattern +E + where namespace(naccumulators=4, indices=array([0, 1, 3, 4]), fundamental=VaspData(array([[ -3.44616154, -19.7728207 , 22.11833881],\n [ -8.54658364, -18.30250658, -15.54715963],\n [ -0.42216143, 8.6986917 , -7.30986158],\n [ -6.76019621, 1.50705188, -1.2590991 ],\n [ 1.09588539, 0.62770383, 2.89870303]])), fundamental_renorm=VaspData(array([[[ 1.31447368, 2.82453526, -2.94655269, 0.31682863,\n 19.32203649, 2.29440329],\n [ 0.93448729, 2.57937955, 1.19506417, 1.30644151,\n 8.56364231, -0.27896706],\n [ 1.03527326, 6.10915995, 9.01056266, -8.2157778 ,\n 5.22714501, 3.18536412]],\n\n [[ 5.48283599, 4.19715567, 4.44685687, 6.43824551,\n -5.52373766, 8.18635823],\n [-13.55073087, 1.22158398, 13.60890006, -16.40893432,\n -26.9191767 , 3.2260027 ],\n [ 28.24687453, 15.48433366, 2.05068837, -15.5690078 ,\n 3.02198295, -4.42156019]],\n\n [[ 0.46848899, 9.64431376, -10.9034597 , 7.77889988,\n -10.92221812, 17.93433268],\n [ -5.97143202, 5.60993696, 8.29856501, -2.96932009,\n 6.38247299, 6.95287228],\n ...} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n\nFundamental gap:\n Temperature \\(K\\) KS gap \\(eV\\) QP gap \\(eV\\) KS-QP gap \\(meV\\)\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}\n^\\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6} \\s*[-+]?\\d*\\.\\d{6}') = .ref +E + and "Electron-phonon bandgap handler with 4 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-10.32472187 -2.04940455 -1.76249068 2.18068674 2.66918125]\n scattering_approx: ['MRTA_TAU' 'SERTA']" = str() +E + and re.MULTILINE = re.MULTILINE +__________________________________ test_len ___________________________________ +tests\calculation\test_electron_phonon_self_energy.py:76: in test_len + assert len(self_energy) == len(self_energy._raw_data.valid_indices) + ^^^^^^^^^^^^^^^^^^^^^ +E AttributeError: 'ElectronPhononSelfEnergy' object has no attribute '_raw_data'. Did you mean: 'from_data'? +_________________________ test_indexing_and_iteration _________________________ +tests\calculation\test_electron_phonon_self_energy.py:84: in test_indexing_and_iteration + assert instance.parent is self_energy +E assert is +E + where = .parent +______________ test_incorrect_selection[invalid_selection=0.01] _______________ +tests\calculation\test_electron_phonon_self_energy.py:189: in test_incorrect_selection + self_energy.select(selection) +src\py4vasp\_calculation\electron_phonon_self_energy.py:268: in select + return self._handler_factory(raw).select(selection) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_self_energy.py:146: in select + indices = self._accumulator().select_indices(selection) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:70: in select_indices + for index_ in self._filter_indices(selection, kwargs_filters) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:82: in _filter_indices + remaining_indices = list(remaining_indices) + ^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:94: in _filter_assignment + if self._match_key_value(index_, key, str(value)): + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:98: in _match_key_value + instance_value = self.get_data(key, index_) + ^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:114: in get_data + self._raise_error_if_not_present(name, expected_name=mu_tag) +src\py4vasp\_calculation\electron_phonon_accumulator.py:122: in _raise_error_if_not_present + valid_names.remove(self._name) +E KeyError: 'electron_phonon_self_energy_handler' +_____________________________ test_print_mapping ______________________________ +tests\calculation\test_electron_phonon_self_energy.py:297: in test_print_mapping + assert re.search(self_energy.ref.mapping_pattern, str(self_energy), re.MULTILINE) +E assert None +E + where None = ('Electron-phonon self energy with 5 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]\n scattering_approx: \\[.*\\]', "Electron-phonon self energy handler with 5 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-17.56267884 -3.48848823 -1.21504499 -0.88958779 15.72487319]\n scattering_approx: ['ERTA_LAMDBA' 'ERTA_TAU' 'MRTA_LAMDBA' 'MRTA_TAU' 'SERTA']", re.MULTILINE) +E + where = re.search +E + and 'Electron-phonon self energy with 5 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]\n scattering_approx: \\[.*\\]' = namespace(eigenvalues=VaspData(array([[[-22.28086492, -13.44773874, 6.37013305],\n [ -3.16280893, 2.81173559, 9.45560547]]])), debye_waller=[VaspData(array([[ 18.50255584, -2.32783758, -2.44522187, 4.18892759,\n 3.00522566, -11.05302037],\n [ 11.1971278 , -4.23930254, 8.08694161, -12.87336232,\n -7.13644385, -15.236876 ],\n [ -7.38825141, -12.71905563, -8.68076598, 7.10879028,\n -3.25210494, 8.57039399],\n [ 22.63794676, 0.03810305, -25.51197157, -4.39087399,\n -12.49321278, -3.46921214]])), VaspData(array([[-16.32009748, -16.43011758, 7.17237632, -25.09938666,\n -1.15079951, -22.15912346],\n [ 17.33743882, 7.85606244, -1.86667992, 8.68090268,\n -19.30600361, 11.67069982],\n [ 21.01959234, -23.25641919, 2.83647841, 4.37003266,\n -3.77920659, 0.11964196],\n [ -1.73235908, -18.2956165 , -24.02048028, -14.15643366,\n -6.37199616, 7.2418473 ]])), VaspData(array([[ 15.31673414, -4.84879213, 13.97888424, -11.84389321,\n 3.01939204, 0.40476462],\n [ 1.53360069, -2.58918443, -6.0906334 , 1.61116058,\n -10.60330148, -2.6566004 ],\n ... [-14.11826492, -17.65226479, 1.84485202, 12.07865246,\n -1.21460353, 5.71945807],\n [-20.46553369, 5.41503234, 5.36529188, 2.16660472,\n 3.90267067, 13.10348344]])), nbands_sum=VaspData(array([ 10, 32, 55, 77, 100], dtype=int32)), selfen_delta=VaspData(array([ -1.21504499, -0.88958779, -3.48848823, -17.56267884,\n 15.72487319])), selfen_carrier_den=array([-5.5779426 , -6.36030311, 19.38344019, -5.5779426 , -6.36030311]), scattering_approx=VaspData(array(['SERTA', 'ERTA_LAMDBA', 'ERTA_TAU', 'MRTA_LAMDBA', 'MRTA_TAU'],\n dtype='.ref +E + and "Electron-phonon self energy handler with 5 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-17.56267884 -3.48848823 -1.21504499 -0.88958779 15.72487319]\n scattering_approx: ['ERTA_LAMDBA' 'ERTA_TAU' 'MRTA_LAMDBA' 'MRTA_TAU' 'SERTA']" = str() +E + and re.MULTILINE = re.MULTILINE +__________________________________ test_len ___________________________________ +tests\calculation\test_electron_phonon_transport.py:94: in test_len + assert len(transport) == len(transport._raw_data.valid_indices) + ^^^^^^^^^^^^^^^^^^^ +E AttributeError: 'ElectronPhononTransport' object has no attribute '_raw_data'. Did you mean: 'from_data'? +_________________________ test_indexing_and_iteration _________________________ +tests\calculation\test_electron_phonon_transport.py:102: in test_indexing_and_iteration + assert instance.parent is transport +E assert is +E + where = .parent +______________ test_incorrect_selection[invalid_selection=0.01] _______________ +tests\calculation\test_electron_phonon_transport.py:240: in test_incorrect_selection + transport.select(selection) +src\py4vasp\_calculation\electron_phonon_transport.py:588: in select + return self._handler_factory(raw).select(selection) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_transport.py:435: in select + return self._select_instances(selection) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_transport.py:438: in _select_instances + indices = self._accumulator().select_indices(selection, *filter_keys) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:70: in select_indices + for index_ in self._filter_indices(selection, kwargs_filters) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:82: in _filter_indices + remaining_indices = list(remaining_indices) + ^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:94: in _filter_assignment + if self._match_key_value(index_, key, str(value)): + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:98: in _match_key_value + instance_value = self.get_data(key, index_) + ^^^^^^^^^^^^^^^^^^^^^^^^^^ +src\py4vasp\_calculation\electron_phonon_accumulator.py:114: in get_data + self._raise_error_if_not_present(name, expected_name=mu_tag) +src\py4vasp\_calculation\electron_phonon_accumulator.py:122: in _raise_error_if_not_present + valid_names.remove(self._name) +E KeyError: 'electron_phonon_transport_handler' +_____________________________ test_print_mapping ______________________________ +tests\calculation\test_electron_phonon_transport.py:475: in test_print_mapping + assert re.search(transport.ref.mapping_pattern, str(transport), re.MULTILINE) +E assert None +E + where None = ('Electron-phonon transport with 5 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]\n scattering_approx: \\[.*\\]', "Electron-phonon transport handler with 5 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-12.92537023 -12.23327944 -0.80108729 2.50646113 4.48223526]\n scattering_approx: ['ERTA_LAMDBA' 'ERTA_TAU' 'MRTA_LAMDBA' 'MRTA_TAU' 'SERTA']", re.MULTILINE) +E + where = re.search +E + and 'Electron-phonon transport with 5 instance\\(s\\):\n selfen_carrier_den: \\[.*\\]\n nbands_sum: \\[.*\\]\n selfen_delta: \\[.*\\]\n scattering_approx: \\[.*\\]' = namespace(temperatures=[array([ 0., 100., 200., 300., 400., 500.]), array([ 0., 100., 200., 300., 400., 500.]), array([ 0., 100., 200., 300., 400., 500.]), array([ 0., 100., 200., 300., 400., 500.]), array([ 0., 100., 200., 300., 400., 500.])], transport_function=VaspData(array([[[[[ 1.15343871e+01, 1.68943639e+00, -7.73719181e+00],\n [ 9.65690947e-01, 1.09984259e+00, -1.94316235e+01],\n [ 5.45903778e+00, -6.53401072e+00, -4.43117265e+00]]],\n\n\n [[[-8.09099478e+00, 4.04558188e+00, -5.29574115e+00],\n [-2.98109677e+00, -1.47448176e+01, 5.66315102e+00],\n [ 4.93091927e+00, -1.64873748e+00, -1.29851841e+01]]],\n\n\n [[[-2.21505214e+00, -4.36085105e+00, -4.13211413e+00],\n [ 1.43926980e+01, 1.74190427e+01, -1.95044529e+01],\n [-4.52561695e-01, 6.31034463e+00, -1.35175447e+01]]],\n\n\n [[[ 1.24624771e+01, 3.17625040e+00, -3.28898994e+00],\n [ 2.05378283e+01, 8.52035417e+00, -5.39064904e+00],\n [-5.46238893e+00, -1.59451253e+01, -2.73101037e+00]]],\n\n\n [[[-3.77560680e+01, -8.23405969e-01, 6.87442062e+00],\n [-4.68548642e+00, 1.35430371e+01, 5.05370359e+00],\n [-9.20683484... , -9.72899234, 13.46365511],\n [-12.64648487, -12.35115495, -1.24292951],\n [ 5.34700525, -3.01601544, -6.29660417]],\n\n [[ 25.71840452, 3.39929278, 33.03263486],\n [ 5.50923261, -1.86665229, -10.14702271],\n [ 3.4960591 , 10.38981741, 5.21138046]],\n\n [[ -3.69250954, 1.98888583, -3.10555872],\n [-18.33541568, 14.65312972, -5.75947149],\n [ 6.6587165 , -11.57366396, 12.92265024]]]])), nbands_sum=VaspData(array([ 10, 32, 55, 77, 100], dtype=int32)), selfen_delta=VaspData(array([ 2.50646113, -0.80108729, -12.23327944, 4.48223526,\n -12.92537023])), selfen_carrier_den=array([-5.5779426 , -6.36030311, 19.38344019, -5.5779426 , -6.36030311]), scattering_approx=VaspData(array(['SERTA', 'ERTA_LAMDBA', 'ERTA_TAU', 'MRTA_LAMDBA', 'MRTA_TAU'],\n dtype='.ref +E + and "Electron-phonon transport handler with 5 instance(s):\n selfen_carrier_den: [-6.36030311 -5.5779426 19.38344019]\n nbands_sum: [ 10 32 55 77 100]\n selfen_delta: [-12.92537023 -12.23327944 -0.80108729 2.50646113 4.48223526]\n scattering_approx: ['ERTA_LAMDBA' 'ERTA_TAU' 'MRTA_LAMDBA' 'MRTA_TAU' 'SERTA']" = str() +E + and re.MULTILINE = re.MULTILINE +=========================== short test summary info =========================== +FAILED tests/calculation/test_electron_phonon_bandgap.py::test_indexing_and_iteration[nonpolarized] +FAILED tests/calculation/test_electron_phonon_bandgap.py::test_indexing_and_iteration[collinear] +FAILED tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[nonpolarized-invalid_selection=0.01] +FAILED tests/calculation/test_electron_phonon_bandgap.py::test_incorrect_selection[collinear-invalid_selection=0.01] +FAILED tests/calculation/test_electron_phonon_bandgap.py::test_print_mapping[nonpolarized] +FAILED tests/calculation/test_electron_phonon_bandgap.py::test_print_mapping[collinear] +FAILED tests/calculation/test_electron_phonon_self_energy.py::test_len - Attr... +FAILED tests/calculation/test_electron_phonon_self_energy.py::test_indexing_and_iteration +FAILED tests/calculation/test_electron_phonon_self_energy.py::test_incorrect_selection[invalid_selection=0.01] +FAILED tests/calculation/test_electron_phonon_self_energy.py::test_print_mapping +FAILED tests/calculation/test_electron_phonon_transport.py::test_len - Attrib... +FAILED tests/calculation/test_electron_phonon_transport.py::test_indexing_and_iteration +FAILED tests/calculation/test_electron_phonon_transport.py::test_incorrect_selection[invalid_selection=0.01] +FAILED tests/calculation/test_electron_phonon_transport.py::test_print_mapping +================== 14 failed, 186 passed, 4 skipped in 6.29s ================== +ERROR conda.cli.main_run:execute(127): `conda run pytest tests/calculation/test_electron_phonon_chemical_potential.py tests/calculation/test_electron_phonon_bandgap.py tests/calculation/test_electron_phonon_self_energy.py tests/calculation/test_electron_phonon_transport.py -v --tb=short` failed. (See above for error) + diff --git a/src/py4vasp/_calculation/_CONTCAR.py b/src/py4vasp/_calculation/_CONTCAR.py index 4c1d29c99..68957c210 100644 --- a/src/py4vasp/_calculation/_CONTCAR.py +++ b/src/py4vasp/_calculation/_CONTCAR.py @@ -1,213 +1,213 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) - -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation.dispatch import ( - DataSource, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import CONTCAR_DB -from py4vasp._third_party import view -from py4vasp._util import check, convert - - -class CONTCARHandler: - """Handler for CONTCAR data — performs all data access and transformation.""" - - def __init__(self, raw_contcar: raw.CONTCAR): - self._raw_contcar = raw_contcar - - @classmethod - def from_data(cls, raw_contcar: raw.CONTCAR) -> "CONTCARHandler": - return cls(raw_contcar) - - def read(self) -> dict: - return self.to_dict() - - def to_dict(self) -> dict: - return { - **self._structure().to_dict(), - "system": convert.text_to_string(self._raw_contcar.system), - **self._read("selective_dynamics"), - **self._read("lattice_velocities"), - **self._read("ion_velocities"), - } - - def to_database(self) -> CONTCAR_DB: - return CONTCAR_DB( - system=( - convert.text_to_string(self._raw_contcar.system) - if not check.is_none(self._raw_contcar.system) - else None - ), - ) - - def to_view(self, supercell=None): - return self._structure().to_view(supercell) - - def __str__(self) -> str: - return "\n".join(self._line_generator()) - - def _line_generator(self): - cell = self._raw_contcar.structure.cell - positions = self._raw_contcar.structure.positions - selective_dynamics = self._raw_contcar.selective_dynamics - yield convert.text_to_string(self._raw_contcar.system) - yield from _cell_lines(cell) - yield self._stoichiometry().to_POSCAR() - if not selective_dynamics.is_none(): - yield "Selective dynamics" - yield "Direct" - yield from _ion_position_lines(positions, selective_dynamics) - yield from _lattice_velocity_lines(self._raw_contcar.lattice_velocities, cell) - yield from _ion_velocity_lines(self._raw_contcar.ion_velocities) - - def _structure(self): - return StructureHandler.from_data(self._raw_contcar.structure) - - def _stoichiometry(self): - return _stoichiometry.Stoichiometry.from_data( - self._raw_contcar.structure.stoichiometry - ) - - def _read(self, key): - data = getattr(self._raw_contcar, key) - return {key: data[:]} if not data.is_none() else {} - - -@quantity("_CONTCAR") -class CONTCAR(view.Mixin): - """CONTCAR contains structural restart-data after a relaxation or MD simulation. - - The CONTCAR contains the final structure of the VASP calculation. It can be used as - input for the next calculation if desired. Depending on the particular setup the - CONTCAR might contain additional information about the system such as the ion - and lattice velocities.""" - - def __init__(self, source, quantity_name="_CONTCAR"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_contcar): - return cls(source=DataSource(raw_contcar)) - - def _handler_factory(self, raw): - return CONTCARHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - CONTCARHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self, selection=None) -> dict: - """Extract the structural data and available additional data to a dictionary.""" - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - CONTCARHandler.read, - ) - - def to_dict(self, selection=None) -> dict: - """Alias for read().""" - return self.read(selection=selection) - - def to_view(self, supercell=None, selection=None): - """Generate a visualization of the final structure.""" - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - CONTCARHandler.to_view, - supercell, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - CONTCARHandler.from_data, - CONTCARHandler.to_database, - ) - - -def _cell_lines(cell): - yield _float_format(_cell_scale(cell.scale), scientific=False).lstrip() - yield from _vectors_to_lines(cell.lattice_vectors) - - -def _cell_scale(scale): - if not scale.is_none(): - return scale[()] - else: - return 1.0 - - -def _ion_position_lines(positions, selective_dynamics): - if selective_dynamics.is_none(): - yield from _vectors_to_lines(positions) - else: - yield from _vectors_and_flags_to_lines(positions, selective_dynamics) - - -def _lattice_velocity_lines(velocities, cell): - if velocities.is_none(): - return - yield "Lattice velocities and vectors" - yield "1" # lattice vectors initialized - yield from _vectors_to_lines(velocities, scientific=True) - lattice_vectors = _cell_scale(cell.scale) * cell.lattice_vectors - yield from _vectors_to_lines(lattice_vectors, scientific=True) - - -def _ion_velocity_lines(velocities): - if velocities.is_none(): - return - yield "Cartesian" - yield from _vectors_to_lines(velocities, scientific=True) - - -def _vectors_to_lines(vectors, scientific=False): - for vector in vectors: - yield _vector_to_line(vector, scientific) - - -def _vectors_and_flags_to_lines(vectors, flags): - for vector, flag in zip(vectors, flags): - yield f"{_vector_to_line(vector, scientific=False)} {_flag_to_line(flag)}" - - -def _vector_to_line(vector, scientific): - insert = "" if scientific else " " - return insert.join(_float_format(x, scientific) for x in vector) - - -def _flag_to_line(flag): - return " ".join(_bool_format(x) for x in flag) - - -def _float_format(number, scientific): - if scientific: - return f"{number:16.8e}" - else: - return f"{number:21.16f}" - - -def _bool_format(value): - return "T" if value else "F" +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import CONTCAR_DB +from py4vasp._third_party import view +from py4vasp._util import check, convert + + +class CONTCARHandler: + """Handler for CONTCAR data — performs all data access and transformation.""" + + def __init__(self, raw_contcar: raw.CONTCAR): + self._raw_contcar = raw_contcar + + @classmethod + def from_data(cls, raw_contcar: raw.CONTCAR) -> "CONTCARHandler": + return cls(raw_contcar) + + def read(self) -> dict: + return self.to_dict() + + def to_dict(self) -> dict: + return { + **self._structure().to_dict(), + "system": convert.text_to_string(self._raw_contcar.system), + **self._read("selective_dynamics"), + **self._read("lattice_velocities"), + **self._read("ion_velocities"), + } + + def to_database(self) -> CONTCAR_DB: + return CONTCAR_DB( + system=( + convert.text_to_string(self._raw_contcar.system) + if not check.is_none(self._raw_contcar.system) + else None + ), + ) + + def to_view(self, supercell=None): + return self._structure().to_view(supercell) + + def __str__(self) -> str: + return "\n".join(self._line_generator()) + + def _line_generator(self): + cell = self._raw_contcar.structure.cell + positions = self._raw_contcar.structure.positions + selective_dynamics = self._raw_contcar.selective_dynamics + yield convert.text_to_string(self._raw_contcar.system) + yield from _cell_lines(cell) + yield self._stoichiometry().to_POSCAR() + if not selective_dynamics.is_none(): + yield "Selective dynamics" + yield "Direct" + yield from _ion_position_lines(positions, selective_dynamics) + yield from _lattice_velocity_lines(self._raw_contcar.lattice_velocities, cell) + yield from _ion_velocity_lines(self._raw_contcar.ion_velocities) + + def _structure(self): + return StructureHandler.from_data(self._raw_contcar.structure) + + def _stoichiometry(self): + return _stoichiometry.Stoichiometry.from_data( + self._raw_contcar.structure.stoichiometry + ) + + def _read(self, key): + data = getattr(self._raw_contcar, key) + return {key: data[:]} if not data.is_none() else {} + + +@quantity("_CONTCAR") +class CONTCAR(view.Mixin): + """CONTCAR contains structural restart-data after a relaxation or MD simulation. + + The CONTCAR contains the final structure of the VASP calculation. It can be used as + input for the next calculation if desired. Depending on the particular setup the + CONTCAR might contain additional information about the system such as the ion + and lattice velocities.""" + + def __init__(self, source, quantity_name="_CONTCAR"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_contcar): + return cls(source=DataSource(raw_contcar)) + + def _handler_factory(self, raw): + return CONTCARHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + CONTCARHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self, selection=None) -> dict: + """Extract the structural data and available additional data to a dictionary.""" + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + CONTCARHandler.read, + ) + + def to_dict(self, selection=None) -> dict: + """Alias for read().""" + return self.read(selection=selection) + + def to_view(self, supercell=None, selection=None): + """Generate a visualization of the final structure.""" + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + CONTCARHandler.to_view, + supercell, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + CONTCARHandler.from_data, + CONTCARHandler.to_database, + ) + + +def _cell_lines(cell): + yield _float_format(_cell_scale(cell.scale), scientific=False).lstrip() + yield from _vectors_to_lines(cell.lattice_vectors) + + +def _cell_scale(scale): + if not scale.is_none(): + return scale[()] + else: + return 1.0 + + +def _ion_position_lines(positions, selective_dynamics): + if selective_dynamics.is_none(): + yield from _vectors_to_lines(positions) + else: + yield from _vectors_and_flags_to_lines(positions, selective_dynamics) + + +def _lattice_velocity_lines(velocities, cell): + if velocities.is_none(): + return + yield "Lattice velocities and vectors" + yield "1" # lattice vectors initialized + yield from _vectors_to_lines(velocities, scientific=True) + lattice_vectors = _cell_scale(cell.scale) * cell.lattice_vectors + yield from _vectors_to_lines(lattice_vectors, scientific=True) + + +def _ion_velocity_lines(velocities): + if velocities.is_none(): + return + yield "Cartesian" + yield from _vectors_to_lines(velocities, scientific=True) + + +def _vectors_to_lines(vectors, scientific=False): + for vector in vectors: + yield _vector_to_line(vector, scientific) + + +def _vectors_and_flags_to_lines(vectors, flags): + for vector, flag in zip(vectors, flags): + yield f"{_vector_to_line(vector, scientific=False)} {_flag_to_line(flag)}" + + +def _vector_to_line(vector, scientific): + insert = "" if scientific else " " + return insert.join(_float_format(x, scientific) for x in vector) + + +def _flag_to_line(flag): + return " ".join(_bool_format(x) for x in flag) + + +def _float_format(number, scientific): + if scientific: + return f"{number:16.8e}" + else: + return f"{number:21.16f}" + + +def _bool_format(value): + return "T" if value else "F" diff --git a/src/py4vasp/_calculation/__init__.py b/src/py4vasp/_calculation/__init__.py index 1ecb83f30..0fdb11e80 100644 --- a/src/py4vasp/_calculation/__init__.py +++ b/src/py4vasp/_calculation/__init__.py @@ -1,5 +1,6 @@ # Copyright © VASP Software GmbH, # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import contextlib import importlib import pathlib from typing import Any, List, Optional, Tuple, Union @@ -8,7 +9,7 @@ from py4vasp._calculation.dispatch import _REGISTRY, FileSource, Group from py4vasp._raw.data import CalculationMetaData, _DatabaseData from py4vasp._raw.data_db import __SCHEMA_VERSION__ -from py4vasp._util import convert, import_ +from py4vasp._util import convert, import_, loadable def _append_database_error( @@ -212,6 +213,112 @@ def path(self): "Return the path in which the calculation is run." return self._path + def selections( + self, method: Optional[str] = None, only_available: bool = False + ) -> dict[str, list[str]]: + """Determine which quantities and selections can be loaded for this calculation. + + This inspects the VASP output files of the calculation and compares them against + the schema defined in :mod:`py4vasp._raw.definition`. For every quantity that + py4vasp can access (e.g. ``"structure"``, ``"band"``, or grouped quantities like + ``"exciton.density"``) it collects all selections (sources) whose data is + present and loadable. + + By default (``only_available=False``), the result reports all schema-defined + selections for each quantity. If *method* is provided, quantities are still + restricted to those implementing that method, but the returned selections are + not filtered by runtime availability. + + When ``only_available=True``, candidate selections are first filtered cheaply + against the schema (only the existence of the relevant datasets is checked). + Every remaining candidate is then confirmed by genuinely invoking the requested + method, so the result only lists selections that truly load. Selections whose + data is present but too large to load raise an out-of-memory error; these are + not reported as loadable. Instead their messages are collected and printed once + at the end, so a single oversized quantity never aborts the inspection. + + Parameters + ---------- + method : str, optional + The method to restrict to. Pass e.g. ``"to_view"`` to only include + quantities that implement visualization. Defaults to ``"read"`` which is + implemented by all quantities. When combined with ``only_available=True``, + the method is also used to confirm runtime loadability. + only_available : bool, optional + If False (default), return all public quantities with their schema-defined + selections. If True, only return quantities for which at least one selection + can be loaded; when *method* is provided, only quantities implementing + that method are included. + + Returns + ------- + dict[str, list[str]] + Maps each quantity call name to a list of selection names (the default + source is reported as ``"default"``). When ``only_available=True``, only + actually loadable selections are included, and quantities with no loadable + selections are omitted. + + Examples + -------- + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + Get all public quantities and their schema-defined selections (default): + + >>> calculation.selections() + {'band': ['default', 'kpoints_opt', 'kpoints_wan'], + 'bandgap': ['default', 'kpoint'], + ...} + + Restrict to quantities implementing a specific method: + + >>> calculation.selections(method="to_view") + {'density': ['default', 'tau'], 'exciton.density': ['default'], ...} + + Get only loadable quantities and their loadable selections: + + >>> calculation.selections(only_available=True) + {'band': ['default', 'kpoints_opt'], 'density': ['default'], ...} + + Combine both options to restrict to loadable quantities implementing a method: + + >>> calculation.selections(method="to_view", only_available=True) + {'density': ['default'], 'exciton.density': ['default'], ...} + """ + _ensure_all_quantities_imported() + result = {} + errors = [] + all_quantities = list(_public_quantities()) + with contextlib.ExitStack() as stack: + open_files = {} + cache = {} + for call_name, schema_name in all_quantities: + if only_available: + sources = loadable.loadable_sources( + self, + call_name, + schema_name, + method, + open_files, + stack, + cache, + QUANTITIES, + errors, + ) + if sources: + result[call_name] = sources + elif method is None or loadable.implements_method( + self, call_name, method + ): + # only_available=False never loads data; it only checks whether + # the quantity implements the requested method. + result[call_name] = loadable.possible_sources(schema_name) + if errors: + header = "Some quantities are available but could not be loaded:" + print("\n".join([header, *errors])) + return dict(sorted(result.items())) + def __getattr__(self, name): # Resolves a quantity (or group) by name from the dispatcher _REGISTRY. Called # only when normal attribute lookup has already failed. @@ -289,6 +396,30 @@ def _compute_database_data(self) -> dict: return properties +def _public_quantities(): + """List (call_name, schema_name) pairs for all user-facing quantities. + + Combines the dispatcher registry (new architecture) with the legacy ``QUANTITIES`` + that are not yet ported. Private quantities (leading underscore) are excluded. + """ + pairs = [] + for key, entry in _REGISTRY.items(): + if key.startswith("_"): + continue + if isinstance(entry, dict): # group of quantities, e.g. exciton.density + for member, dispatcher_cls in entry.items(): + if member.startswith("_"): + continue + pairs.append((f"{key}.{member}", dispatcher_cls._quantity_name)) + else: + pairs.append((key, entry._quantity_name)) + for quantity in QUANTITIES: + if quantity.startswith("_") or quantity in _REGISTRY: + continue + pairs.append((quantity, quantity)) + return pairs + + def _ensure_all_quantities_imported(): """Import all quantity modules so that _REGISTRY is fully populated.""" calc_pkg = importlib.import_module("py4vasp._calculation") diff --git a/src/py4vasp/_calculation/_dispersion.py b/src/py4vasp/_calculation/_dispersion.py index b9df67877..3c792814b 100644 --- a/src/py4vasp/_calculation/_dispersion.py +++ b/src/py4vasp/_calculation/_dispersion.py @@ -1,224 +1,223 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import numpy as np - -import py4vasp._third_party.graph as _graph -from py4vasp._calculation import projector -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.kpoint import KpointHandler -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Dispersion_DB -from py4vasp._util import check - - -class DispersionHandler: - """Handler for dispersion (band/phonon) data.""" - - def __init__(self, raw_dispersion: raw.Dispersion): - self._raw_dispersion = raw_dispersion - - @classmethod - def from_data(cls, raw_dispersion: raw.Dispersion) -> "DispersionHandler": - return cls(raw_dispersion) - - def __str__(self): - return f"""band data: - {self._kpoints().number_kpoints()} k-points - {self._raw_dispersion.eigenvalues.shape[-1]} bands""" - - def read(self) -> dict: - return self.to_dict() - - def to_dict(self) -> dict: - kpoints = self._kpoints() - kpoint_labels = kpoints.labels() - labels_dict = {} if kpoint_labels is None else {"kpoint_labels": kpoint_labels} - return { - "kpoint_distances": kpoints.distances(), - **labels_dict, - "eigenvalues": self._raw_dispersion.eigenvalues[:], - } - - def to_database(self) -> dict: - eigenvalues = ( - self._raw_dispersion.eigenvalues[:] - if not check.is_none(self._raw_dispersion.eigenvalues) - else None - ) - min_eigenvalue = float(np.min(eigenvalues)) if eigenvalues is not None else None - max_eigenvalue = float(np.max(eigenvalues)) if eigenvalues is not None else None - - min_eigenvalue_up, max_eigenvalue_up = None, None - min_eigenvalue_down, max_eigenvalue_down = None, None - if self._spin_polarized(): - eigenvalues_up = eigenvalues[0] - eigenvalues_down = eigenvalues[1] - min_eigenvalue_up = float(np.min(eigenvalues_up)) - max_eigenvalue_up = float(np.max(eigenvalues_up)) - min_eigenvalue_down = float(np.min(eigenvalues_down)) - max_eigenvalue_down = float(np.max(eigenvalues_down)) - - return Dispersion_DB( - eigenvalue_min=min_eigenvalue, - eigenvalue_max=max_eigenvalue, - eigenvalue_min_up=min_eigenvalue_up, - eigenvalue_max_up=max_eigenvalue_up, - eigenvalue_min_down=min_eigenvalue_down, - eigenvalue_max_down=max_eigenvalue_down, - ) - - def plot(self, projections=None): - data = self.to_dict() - projections = self._use_projections_or_default(projections) - return _graph.Graph( - series=_band_structure(data, projections), - xticks=_xticks(data, self._kpoints().line_length()), - ) - - def _kpoints(self): - return KpointHandler.from_data(self._raw_dispersion.kpoints) - - def _use_projections_or_default(self, projections): - if projections is not None: - return projections - elif self._spin_polarized(): - return {"up": None, "down": None} - else: - return {"bands": None} - - def _spin_polarized(self): - eigenvalues = self._raw_dispersion.eigenvalues - return eigenvalues.ndim == 3 and eigenvalues.shape[0] == 2 - - -@quantity("_dispersion") -class Dispersion: - """Generic class for all dispersions (electrons, phonons).""" - - def __init__(self, source, quantity_name="dispersion"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_dispersion): - return cls(source=DataSource(raw_dispersion)) - - def _handler_factory(self, raw): - return DispersionHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DispersionHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self, selection=None) -> dict: - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DispersionHandler.read, - ) - - def to_dict(self, selection=None) -> dict: - return self.read(selection=selection) - - def plot(self, projections=None): - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - DispersionHandler.plot, - projections, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - DispersionHandler.from_data, - DispersionHandler.to_database, - ) - - -def _band_structure(data, projections): - spin_projections = projections.get(projector.SPIN_PROJECTION, []) - return [ - _make_series(data, label, weight, label in spin_projections) - for label, weight in projections.items() - if label != projector.SPIN_PROJECTION - ] - - -def _make_series(data, label, weight, is_spin_projection): - options = {} - if not check.is_none(weight): - options["weight"] = weight.T - if is_spin_projection: - options["marker"] = "o" - options["weight_mode"] = "color" - x = data["kpoint_distances"] - y = _get_bands(data["eigenvalues"], label) - return _graph.Series(x, y, label, **options) - - -def _get_bands(eigenvalues, label): - if eigenvalues.ndim == 2: - return eigenvalues.T - elif "down" in label: - return eigenvalues[1].T - else: - return eigenvalues[0].T - - -def _xticks(data, line_length): - ticks, labels = _degenerate_ticks_and_labels(data, line_length) - return _filter_unique(ticks, labels) - - -def _degenerate_ticks_and_labels(data, line_length): - labels = _tick_labels(data) - mask = _use_labels_and_line_edges(labels, line_length) - return data["kpoint_distances"][mask], labels[mask] - - -def _tick_labels(data): - kpoint_labels = data.get("kpoint_labels") - if kpoint_labels is None: - return np.zeros(len(data["kpoint_distances"]), str) - else: - return np.array(kpoint_labels) - - -def _use_labels_and_line_edges(labels, line_length): - mask = labels != "" - mask[::line_length] = True - mask[-1] = True - return mask - - -def _filter_unique(ticks, labels): - result = {} - for tick, label in zip(ticks, labels): - if tick in result: - previous_label = result[tick] - if previous_label != "" and previous_label != label: - label = previous_label + "|" + label - result[tick] = label - return result +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import numpy as np + +import py4vasp._third_party.graph as _graph +from py4vasp._calculation import projector +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.kpoint import KpointHandler +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Dispersion_DB +from py4vasp._util import check + + +class DispersionHandler: + """Handler for dispersion (band/phonon) data.""" + + def __init__(self, raw_dispersion: raw.Dispersion): + self._raw_dispersion = raw_dispersion + + @classmethod + def from_data(cls, raw_dispersion: raw.Dispersion) -> "DispersionHandler": + return cls(raw_dispersion) + + def __str__(self): + return f"""band data: + {self._kpoints().number_kpoints()} k-points + {self._raw_dispersion.eigenvalues.shape[-1]} bands""" + + def read(self) -> dict: + return self.to_dict() + + def to_dict(self) -> dict: + kpoints = self._kpoints() + kpoint_labels = kpoints.labels() + labels_dict = {} if kpoint_labels is None else {"kpoint_labels": kpoint_labels} + return { + "kpoint_distances": kpoints.distances(), + **labels_dict, + "eigenvalues": self._raw_dispersion.eigenvalues[:], + } + + def to_database(self) -> dict: + eigenvalues = ( + self._raw_dispersion.eigenvalues[:] + if not check.is_none(self._raw_dispersion.eigenvalues) + else None + ) + min_eigenvalue = float(np.min(eigenvalues)) if eigenvalues is not None else None + max_eigenvalue = float(np.max(eigenvalues)) if eigenvalues is not None else None + + min_eigenvalue_up, max_eigenvalue_up = None, None + min_eigenvalue_down, max_eigenvalue_down = None, None + if self._spin_polarized(): + eigenvalues_up = eigenvalues[0] + eigenvalues_down = eigenvalues[1] + min_eigenvalue_up = float(np.min(eigenvalues_up)) + max_eigenvalue_up = float(np.max(eigenvalues_up)) + min_eigenvalue_down = float(np.min(eigenvalues_down)) + max_eigenvalue_down = float(np.max(eigenvalues_down)) + + return Dispersion_DB( + eigenvalue_min=min_eigenvalue, + eigenvalue_max=max_eigenvalue, + eigenvalue_min_up=min_eigenvalue_up, + eigenvalue_max_up=max_eigenvalue_up, + eigenvalue_min_down=min_eigenvalue_down, + eigenvalue_max_down=max_eigenvalue_down, + ) + + def plot(self, projections=None): + data = self.to_dict() + projections = self._use_projections_or_default(projections) + return _graph.Graph( + series=_band_structure(data, projections), + xticks=_xticks(data, self._kpoints().line_length()), + ) + + def _kpoints(self): + return KpointHandler.from_data(self._raw_dispersion.kpoints) + + def _use_projections_or_default(self, projections): + if projections is not None: + return projections + elif self._spin_polarized(): + return {"up": None, "down": None} + else: + return {"bands": None} + + def _spin_polarized(self): + eigenvalues = self._raw_dispersion.eigenvalues + return eigenvalues.ndim == 3 and eigenvalues.shape[0] == 2 + + +@quantity("_dispersion") +class Dispersion: + """Generic class for all dispersions (electrons, phonons).""" + + def __init__(self, source, quantity_name="dispersion"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_dispersion): + return cls(source=DataSource(raw_dispersion)) + + def _handler_factory(self, raw): + return DispersionHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DispersionHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self, selection=None) -> dict: + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DispersionHandler.read, + ) + + def to_dict(self, selection=None) -> dict: + return self.read(selection=selection) + + def plot(self, projections=None): + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + DispersionHandler.plot, + projections, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + DispersionHandler.from_data, + DispersionHandler.to_database, + ) + + +def _band_structure(data, projections): + spin_projections = projections.get(projector.SPIN_PROJECTION, []) + return [ + _make_series(data, label, weight, label in spin_projections) + for label, weight in projections.items() + if label != projector.SPIN_PROJECTION + ] + + +def _make_series(data, label, weight, is_spin_projection): + options = {} + if not check.is_none(weight): + options["weight"] = weight.T + if is_spin_projection: + options["marker"] = "o" + options["weight_mode"] = "color" + x = data["kpoint_distances"] + y = _get_bands(data["eigenvalues"], label) + return _graph.Series(x, y, label, **options) + + +def _get_bands(eigenvalues, label): + if eigenvalues.ndim == 2: + return eigenvalues.T + elif "down" in label: + return eigenvalues[1].T + else: + return eigenvalues[0].T + + +def _xticks(data, line_length): + ticks, labels = _degenerate_ticks_and_labels(data, line_length) + return _filter_unique(ticks, labels) + + +def _degenerate_ticks_and_labels(data, line_length): + labels = _tick_labels(data) + mask = _use_labels_and_line_edges(labels, line_length) + return data["kpoint_distances"][mask], labels[mask] + + +def _tick_labels(data): + kpoint_labels = data.get("kpoint_labels") + if kpoint_labels is None: + return np.zeros(len(data["kpoint_distances"]), str) + else: + return np.array(kpoint_labels) + + +def _use_labels_and_line_edges(labels, line_length): + mask = labels != "" + mask[::line_length] = True + mask[-1] = True + return mask + + +def _filter_unique(ticks, labels): + result = {} + for tick, label in zip(ticks, labels): + if tick in result: + previous_label = result[tick] + if previous_label != "" and previous_label != label: + label = previous_label + "|" + label + result[tick] = label + return result diff --git a/src/py4vasp/_calculation/_stoichiometry.py b/src/py4vasp/_calculation/_stoichiometry.py index 53564018f..726ca1646 100644 --- a/src/py4vasp/_calculation/_stoichiometry.py +++ b/src/py4vasp/_calculation/_stoichiometry.py @@ -1,456 +1,456 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import itertools - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation.dispatch import ( - DataSource, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.selection import Selection -from py4vasp._raw.data_db import Stoichiometry_DB -from py4vasp._util import check, convert, database, documentation, import_, select - -mdtraj = import_.optional("mdtraj") -pd = import_.optional("pandas") - -_subscript = "_" - - -ion_types_documentation = """\ -ion_types : Sequence - Overwrite the ion types present in the raw data.""" - - -class StoichiometryHandler: - """Processes stoichiometry data from a single raw.Stoichiometry object.""" - - def __init__(self, raw_stoichiometry: raw.Stoichiometry): - self._raw_stoichiometry = raw_stoichiometry - - @classmethod - def from_data(cls, raw_stoichiometry: raw.Stoichiometry) -> "StoichiometryHandler": - return cls(raw_stoichiometry) - - def read(self, ion_types=None) -> dict: - """Read the stoichiometry into a dictionary of selections.""" - return {**self._default_selection(), **self._specific_selection(ion_types)} - - def to_dict(self, ion_types=None) -> dict: - return self.read(ion_types=ion_types) - - def to_string(self, ion_types=None) -> str: - """Convert the stoichiometry into a string.""" - number_suffix = lambda number: str(number) if number > 1 else "" - dummy_type = lambda index: f"({chr(ord('A') + index)})" - return self._create_repr(number_suffix, dummy_type, ion_types) - - def __str__(self): - return self.to_string() - - def to_html(self) -> str: - """HTML representation of the stoichiometry.""" - number_suffix = lambda number: f"{number}" if number > 1 else "" - dummy_type = lambda index: f"{chr(ord('A') + index)}" - return self._create_repr(number_suffix, dummy_type) - - def to_frame(self, ion_types=None): - """Convert the stoichiometry to a DataFrame.""" - return pd.DataFrame( - {"name": self.names(ion_types), "element": self.elements(ion_types)} - ) - - def to_mdtraj(self, ion_types=None): - """Convert the stoichiometry to a mdtraj.Topology.""" - df = self.to_frame(ion_types) - df["serial"] = None - df["resSeq"] = 0 - df["resName"] = "crystal" - df["chainID"] = 0 - df["formal_charge"] = 0 - return mdtraj.Topology.from_dataframe(df) - - def to_POSCAR(self, format_newline="", ion_types=None) -> str: - """Generate the stoichiometry lines for the POSCAR file.""" - error_message = "The formatting information must be a string." - check.raise_error_if_not_string(format_newline, error_message) - number_ion_types = " ".join( - str(x) for x in self._raw_stoichiometry.number_ion_types - ) - if ion_types is None and check.is_none(self._raw_stoichiometry.ion_types): - return number_ion_types - else: - ion_types_str = " ".join(self._ion_types(ion_types)) - return ion_types_str + format_newline + "\n" + number_ion_types - - def names(self, ion_types=None) -> list: - """Extract the labels of all atoms.""" - atom_dict = self.read(ion_types) - return [val.label for val in atom_dict.values() if _subscript in val.label] - - def elements(self, ion_types=None) -> list: - """Extract the element of all atoms.""" - repeated_types = (itertools.repeat(*x) for x in self._type_numbers(ion_types)) - return list(itertools.chain.from_iterable(repeated_types)) - - def ion_types_list(self, ion_types=None) -> list: - """Return the type of all ions in the system as string.""" - return list(dict.fromkeys(self._ion_types(ion_types))) - - def number_atoms(self) -> int: - """Return the number of atoms in the system.""" - return int(np.sum(self._raw_stoichiometry.number_ion_types)) - - def to_database(self) -> Stoichiometry_DB: - """Return database-ready stoichiometry data.""" - ion_types = ( - list(self._ion_types(None)) - if not check.is_none(self._raw_stoichiometry.ion_types) - else None - ) - num_ion_types = ( - list(self._raw_stoichiometry.number_ion_types) - if not check.is_none(self._raw_stoichiometry.number_ion_types) - else None - ) - formula, compound, simple_types, simple_numbers, primitive_numbers = ( - database.get_formula_and_compound(ion_types, num_ion_types) - ) - return Stoichiometry_DB( - ion_types=simple_types, - num_ion_types=simple_numbers, - num_ion_types_primitive=primitive_numbers, - formula=formula, - compound=compound, - ) - - def _create_repr(self, number_suffix, dummy_type, ion_types=None): - ion_string = lambda ion, number: f"{ion}{number_suffix(number)}" - if ion_types is None and check.is_none(self._raw_stoichiometry.ion_types): - number_ion_types = range(len(self._raw_stoichiometry.number_ion_types)) - ion_types = [dummy_type(i) for i in number_ion_types] - elif ion_types is None: - ion_types = self._raw_stoichiometry.ion_types - total_type_numbers = self._total_type_numbers(ion_types) - return "".join(ion_string(*item) for item in total_type_numbers.items()) - - def _total_type_numbers(self, ion_types): - result = {} - for ion_type, number in self._type_numbers(ion_types): - result.setdefault(ion_type, 0) - result[ion_type] += number - return result - - def _default_selection(self): - num_atoms = self.number_atoms() - return {select.all: Selection(indices=slice(0, num_atoms))} - - def _specific_selection(self, ion_types): - result = {} - for i, element in enumerate(self.elements(ion_types)): - result.setdefault(element, Selection(indices=[], label=element)) - result[element].indices.append(i) - label = f"{element}{_subscript}{len(result[element].indices)}" - result[str(i + 1)] = Selection(indices=[i], label=label) - return _merge_to_slice_if_possible(result) - - def _type_numbers(self, ion_types): - return zip(self._ion_types(ion_types), self._raw_stoichiometry.number_ion_types) - - def _ion_types(self, ion_types): - ion_types = ( - self._raw_stoichiometry.ion_types if ion_types is None else ion_types - ) - if check.is_none(ion_types): - message = "If the ion types are not defined, you must pass them as argument to the function." - raise exception.IncorrectUsage(message) - clean_string = lambda ion_type: convert.text_to_string(ion_type).strip() - return (clean_string(ion_type) for ion_type in ion_types) - - -@quantity("_stoichiometry") -class Stoichiometry: - """The stoichiometry of the crystal describes how many ions of each type exist in a crystal.""" - - def __init__(self, source, quantity_name="stoichiometry"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_stoichiometry): - return cls(source=DataSource(raw_stoichiometry)) - - @classmethod - def from_ase(cls, structure): - """Generate a stoichiometry from the given ase Atoms object.""" - return cls.from_data(raw_stoichiometry_from_ase(structure)) - - def _handler_factory(self, raw): - return StoichiometryHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - StoichiometryHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def _repr_html_(self): - return merge_strings( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.to_html, - ) - - @documentation.format(ion_types=ion_types_documentation) - def to_string(self, ion_types=None): - """Convert the stoichiometry into a string. - - This method is equivalent to calling str() on the stoichiometry except that it - allows to overwrite the ion types. - - Parameters - ---------- - {ion_types} - - Returns - ------- - str - String representation of the stoichiometry. - """ - return merge_strings( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.to_string, - ion_types, - ) - - def read(self, ion_types=None): - "Convenient wrapper around to_dict. Check that function for examples and optional arguments." - return self.to_dict(ion_types=ion_types) - - @documentation.format(ion_types=ion_types_documentation) - def to_dict(self, ion_types=None): - """Read the stoichiometry and convert it to a dictionary. - - Parameters - ---------- - {ion_types} - - Returns - ------- - dict - A map from particular labels to the corresponding atom indices. For - every atom a single label of the form *Element*_*Number* (e.g. Sr_1) - is constructed. In addition there is a map from atoms of a specific - element type to all indices of that element and from atom-indices - strings to atom-indices integers. These access strings are used - throughout all of the refinement classes. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.to_dict, - ion_types, - ) - - @documentation.format(ion_types=ion_types_documentation) - def to_frame(self, ion_types=None): - """Convert the stoichiometry to a DataFrame - - Parameters - ---------- - {ion_types} - - Returns - ------- - pd.DataFrame - The dataframe matches atom label and element type. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.to_frame, - ion_types, - ) - - @documentation.format(ion_types=ion_types_documentation) - def to_mdtraj(self, ion_types=None): - """Convert the stoichiometry to a mdtraj.Topology. - - Parameters - ---------- - {ion_types} - - Returns - ------- - mdtraj.Topology - Converts the stoichiometry into an object that can be used for mdtraj. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.to_mdtraj, - ion_types, - ) - - @documentation.format(ion_types=ion_types_documentation) - def to_POSCAR(self, format_newline="", ion_types=None): - """Generate the stoichiometry lines for the POSCAR file. - - Parameters - ---------- - format_newline : str - If you want to display the POSCAR file in a particular way, you can - use an additional string to add formatting. - {ion_types} - - Returns - ------- - str - A string used to describe the atoms in the system in the POSCAR file - augmented by the additional formatting string, if given. - """ - return merge_strings( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.to_POSCAR, - format_newline, - ion_types, - ) - - @documentation.format(ion_types=ion_types_documentation) - def names(self, ion_types=None): - """Extract the labels of all atoms. - - Parameters - ---------- - {ion_types} - - Returns - ------- - list - List of unique string labeling each ion. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.names, - ion_types, - ) - - @documentation.format(ion_types=ion_types_documentation) - def elements(self, ion_types=None): - """Extract the element of all atoms. - - Parameters - ---------- - {ion_types} - - Returns - ------- - list - List of strings specifying the element of each ion. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.elements, - ion_types, - ) - - @documentation.format(ion_types=ion_types_documentation) - def ion_types(self, ion_types=None): - """Return the type of all ions in the system as string. - - Parameters - ---------- - {ion_types} - - Returns - ------- - list - List of unique elements. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.ion_types_list, - ion_types, - ) - - def number_atoms(self): - "Return the number of atoms in the system." - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StoichiometryHandler.number_atoms, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - StoichiometryHandler.from_data, - StoichiometryHandler.to_database, - ) - - -def raw_stoichiometry_from_ase(structure): - """Convert the given ase Atoms object to a raw.Stoichiometry.""" - number_ion_types = [] - ion_types = [] - for element in structure.symbols: - if ion_types and ion_types[-1] == element: - number_ion_types[-1] += 1 - else: - ion_types.append(element) - number_ion_types.append(1) - return raw.Stoichiometry(number_ion_types, ion_types) - - -def _merge_to_slice_if_possible(selections): - for selection in selections.values(): - if _is_slice(selection.indices): - selection.indices = _to_slice(selection.indices) - return selections - - -def _is_slice(indices): - assert sorted(indices) == indices - return len(indices) == indices[-1] - indices[0] + 1 - - -def _to_slice(indices): - return slice(indices[0], indices[-1] + 1) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import itertools + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.selection import Selection +from py4vasp._raw.data_db import Stoichiometry_DB +from py4vasp._util import check, convert, database, documentation, import_, select + +mdtraj = import_.optional("mdtraj") +pd = import_.optional("pandas") + +_subscript = "_" + + +ion_types_documentation = """\ +ion_types : Sequence + Overwrite the ion types present in the raw data.""" + + +class StoichiometryHandler: + """Processes stoichiometry data from a single raw.Stoichiometry object.""" + + def __init__(self, raw_stoichiometry: raw.Stoichiometry): + self._raw_stoichiometry = raw_stoichiometry + + @classmethod + def from_data(cls, raw_stoichiometry: raw.Stoichiometry) -> "StoichiometryHandler": + return cls(raw_stoichiometry) + + def read(self, ion_types=None) -> dict: + """Read the stoichiometry into a dictionary of selections.""" + return {**self._default_selection(), **self._specific_selection(ion_types)} + + def to_dict(self, ion_types=None) -> dict: + return self.read(ion_types=ion_types) + + def to_string(self, ion_types=None) -> str: + """Convert the stoichiometry into a string.""" + number_suffix = lambda number: str(number) if number > 1 else "" + dummy_type = lambda index: f"({chr(ord('A') + index)})" + return self._create_repr(number_suffix, dummy_type, ion_types) + + def __str__(self): + return self.to_string() + + def to_html(self) -> str: + """HTML representation of the stoichiometry.""" + number_suffix = lambda number: f"{number}" if number > 1 else "" + dummy_type = lambda index: f"{chr(ord('A') + index)}" + return self._create_repr(number_suffix, dummy_type) + + def to_frame(self, ion_types=None): + """Convert the stoichiometry to a DataFrame.""" + return pd.DataFrame( + {"name": self.names(ion_types), "element": self.elements(ion_types)} + ) + + def to_mdtraj(self, ion_types=None): + """Convert the stoichiometry to a mdtraj.Topology.""" + df = self.to_frame(ion_types) + df["serial"] = None + df["resSeq"] = 0 + df["resName"] = "crystal" + df["chainID"] = 0 + df["formal_charge"] = 0 + return mdtraj.Topology.from_dataframe(df) + + def to_POSCAR(self, format_newline="", ion_types=None) -> str: + """Generate the stoichiometry lines for the POSCAR file.""" + error_message = "The formatting information must be a string." + check.raise_error_if_not_string(format_newline, error_message) + number_ion_types = " ".join( + str(x) for x in self._raw_stoichiometry.number_ion_types + ) + if ion_types is None and check.is_none(self._raw_stoichiometry.ion_types): + return number_ion_types + else: + ion_types_str = " ".join(self._ion_types(ion_types)) + return ion_types_str + format_newline + "\n" + number_ion_types + + def names(self, ion_types=None) -> list: + """Extract the labels of all atoms.""" + atom_dict = self.read(ion_types) + return [val.label for val in atom_dict.values() if _subscript in val.label] + + def elements(self, ion_types=None) -> list: + """Extract the element of all atoms.""" + repeated_types = (itertools.repeat(*x) for x in self._type_numbers(ion_types)) + return list(itertools.chain.from_iterable(repeated_types)) + + def ion_types_list(self, ion_types=None) -> list: + """Return the type of all ions in the system as string.""" + return list(dict.fromkeys(self._ion_types(ion_types))) + + def number_atoms(self) -> int: + """Return the number of atoms in the system.""" + return int(np.sum(self._raw_stoichiometry.number_ion_types)) + + def to_database(self) -> Stoichiometry_DB: + """Return database-ready stoichiometry data.""" + ion_types = ( + list(self._ion_types(None)) + if not check.is_none(self._raw_stoichiometry.ion_types) + else None + ) + num_ion_types = ( + list(self._raw_stoichiometry.number_ion_types) + if not check.is_none(self._raw_stoichiometry.number_ion_types) + else None + ) + formula, compound, simple_types, simple_numbers, primitive_numbers = ( + database.get_formula_and_compound(ion_types, num_ion_types) + ) + return Stoichiometry_DB( + ion_types=simple_types, + num_ion_types=simple_numbers, + num_ion_types_primitive=primitive_numbers, + formula=formula, + compound=compound, + ) + + def _create_repr(self, number_suffix, dummy_type, ion_types=None): + ion_string = lambda ion, number: f"{ion}{number_suffix(number)}" + if ion_types is None and check.is_none(self._raw_stoichiometry.ion_types): + number_ion_types = range(len(self._raw_stoichiometry.number_ion_types)) + ion_types = [dummy_type(i) for i in number_ion_types] + elif ion_types is None: + ion_types = self._raw_stoichiometry.ion_types + total_type_numbers = self._total_type_numbers(ion_types) + return "".join(ion_string(*item) for item in total_type_numbers.items()) + + def _total_type_numbers(self, ion_types): + result = {} + for ion_type, number in self._type_numbers(ion_types): + result.setdefault(ion_type, 0) + result[ion_type] += number + return result + + def _default_selection(self): + num_atoms = self.number_atoms() + return {select.all: Selection(indices=slice(0, num_atoms))} + + def _specific_selection(self, ion_types): + result = {} + for i, element in enumerate(self.elements(ion_types)): + result.setdefault(element, Selection(indices=[], label=element)) + result[element].indices.append(i) + label = f"{element}{_subscript}{len(result[element].indices)}" + result[str(i + 1)] = Selection(indices=[i], label=label) + return _merge_to_slice_if_possible(result) + + def _type_numbers(self, ion_types): + return zip(self._ion_types(ion_types), self._raw_stoichiometry.number_ion_types) + + def _ion_types(self, ion_types): + ion_types = ( + self._raw_stoichiometry.ion_types if ion_types is None else ion_types + ) + if check.is_none(ion_types): + message = "If the ion types are not defined, you must pass them as argument to the function." + raise exception.IncorrectUsage(message) + clean_string = lambda ion_type: convert.text_to_string(ion_type).strip() + return (clean_string(ion_type) for ion_type in ion_types) + + +@quantity("_stoichiometry") +class Stoichiometry: + """The stoichiometry of the crystal describes how many ions of each type exist in a crystal.""" + + def __init__(self, source, quantity_name="stoichiometry"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_stoichiometry): + return cls(source=DataSource(raw_stoichiometry)) + + @classmethod + def from_ase(cls, structure): + """Generate a stoichiometry from the given ase Atoms object.""" + return cls.from_data(raw_stoichiometry_from_ase(structure)) + + def _handler_factory(self, raw): + return StoichiometryHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + StoichiometryHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def _repr_html_(self): + return merge_strings( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.to_html, + ) + + @documentation.format(ion_types=ion_types_documentation) + def to_string(self, ion_types=None): + """Convert the stoichiometry into a string. + + This method is equivalent to calling str() on the stoichiometry except that it + allows to overwrite the ion types. + + Parameters + ---------- + {ion_types} + + Returns + ------- + str + String representation of the stoichiometry. + """ + return merge_strings( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.to_string, + ion_types, + ) + + def read(self, ion_types=None): + "Convenient wrapper around to_dict. Check that function for examples and optional arguments." + return self.to_dict(ion_types=ion_types) + + @documentation.format(ion_types=ion_types_documentation) + def to_dict(self, ion_types=None): + """Read the stoichiometry and convert it to a dictionary. + + Parameters + ---------- + {ion_types} + + Returns + ------- + dict + A map from particular labels to the corresponding atom indices. For + every atom a single label of the form *Element*_*Number* (e.g. Sr_1) + is constructed. In addition there is a map from atoms of a specific + element type to all indices of that element and from atom-indices + strings to atom-indices integers. These access strings are used + throughout all of the refinement classes. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.to_dict, + ion_types, + ) + + @documentation.format(ion_types=ion_types_documentation) + def to_frame(self, ion_types=None): + """Convert the stoichiometry to a DataFrame + + Parameters + ---------- + {ion_types} + + Returns + ------- + pd.DataFrame + The dataframe matches atom label and element type. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.to_frame, + ion_types, + ) + + @documentation.format(ion_types=ion_types_documentation) + def to_mdtraj(self, ion_types=None): + """Convert the stoichiometry to a mdtraj.Topology. + + Parameters + ---------- + {ion_types} + + Returns + ------- + mdtraj.Topology + Converts the stoichiometry into an object that can be used for mdtraj. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.to_mdtraj, + ion_types, + ) + + @documentation.format(ion_types=ion_types_documentation) + def to_POSCAR(self, format_newline="", ion_types=None): + """Generate the stoichiometry lines for the POSCAR file. + + Parameters + ---------- + format_newline : str + If you want to display the POSCAR file in a particular way, you can + use an additional string to add formatting. + {ion_types} + + Returns + ------- + str + A string used to describe the atoms in the system in the POSCAR file + augmented by the additional formatting string, if given. + """ + return merge_strings( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.to_POSCAR, + format_newline, + ion_types, + ) + + @documentation.format(ion_types=ion_types_documentation) + def names(self, ion_types=None): + """Extract the labels of all atoms. + + Parameters + ---------- + {ion_types} + + Returns + ------- + list + List of unique string labeling each ion. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.names, + ion_types, + ) + + @documentation.format(ion_types=ion_types_documentation) + def elements(self, ion_types=None): + """Extract the element of all atoms. + + Parameters + ---------- + {ion_types} + + Returns + ------- + list + List of strings specifying the element of each ion. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.elements, + ion_types, + ) + + @documentation.format(ion_types=ion_types_documentation) + def ion_types(self, ion_types=None): + """Return the type of all ions in the system as string. + + Parameters + ---------- + {ion_types} + + Returns + ------- + list + List of unique elements. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.ion_types_list, + ion_types, + ) + + def number_atoms(self): + "Return the number of atoms in the system." + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StoichiometryHandler.number_atoms, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + StoichiometryHandler.from_data, + StoichiometryHandler.to_database, + ) + + +def raw_stoichiometry_from_ase(structure): + """Convert the given ase Atoms object to a raw.Stoichiometry.""" + number_ion_types = [] + ion_types = [] + for element in structure.symbols: + if ion_types and ion_types[-1] == element: + number_ion_types[-1] += 1 + else: + ion_types.append(element) + number_ion_types.append(1) + return raw.Stoichiometry(number_ion_types, ion_types) + + +def _merge_to_slice_if_possible(selections): + for selection in selections.values(): + if _is_slice(selection.indices): + selection.indices = _to_slice(selection.indices) + return selections + + +def _is_slice(indices): + assert sorted(indices) == indices + return len(indices) == indices[-1] - indices[0] + 1 + + +def _to_slice(indices): + return slice(indices[0], indices[-1] + 1) diff --git a/src/py4vasp/_calculation/band.py b/src/py4vasp/_calculation/band.py index 746eb7d56..85afbab43 100644 --- a/src/py4vasp/_calculation/band.py +++ b/src/py4vasp/_calculation/band.py @@ -1,881 +1,879 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from __future__ import annotations - -import pathlib -from typing import Any, Iterable, List, Optional - -import numpy as np -from numpy.typing import ArrayLike - -from py4vasp import exception -from py4vasp._calculation import projector -from py4vasp._calculation._dispersion import DispersionHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.kpoint import KpointHandler -from py4vasp._calculation.projector import ProjectorHandler -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Band_DB -from py4vasp._third_party import graph -from py4vasp._util import ( - check, - documentation, - import_, - index, - select, - slicing, -) - -pd = import_.optional("pandas") -pretty = import_.optional("IPython.lib.pretty") - -_OCCUPATION_CUTOFF = 1e-2 - -# Cartesian axis (x=0, y=1, z=2) corresponding to each accepted spin-component token. -_SPIN_TOKEN_TO_CARTESIAN_AXIS = { - "x": 0, - "sigma_x": 0, - "sigma_1": 0, - "y": 1, - "sigma_y": 1, - "sigma_2": 1, - "z": 2, - "sigma_z": 2, - "sigma_3": 2, -} - - -class BandHandler: - """Handler for electronic band structure data.""" - - def __init__(self, raw_band: raw.Band): - self._raw_band = raw_band - - @classmethod - def from_data(cls, raw_band: raw.Band) -> "BandHandler": - return cls(raw_band) - - def __str__(self) -> str: - return f""" -{"spin polarized" if self._is_collinear() else ""} band data: - {self._raw_band.dispersion.eigenvalues.shape[1]} k-points - {self._raw_band.dispersion.eigenvalues.shape[2]} bands -{str(self._projector())} - """.strip() - - def to_dict(self, selection=None, fermi_energy=None) -> dict[str, Any]: - dispersion = self._dispersion().to_dict() - eigenvalues = dispersion.pop("eigenvalues") - return { - **dispersion, - "fermi_energy": self._raw_band.fermi_energy, - **self._shift_dispersion_by_fermi_energy(eigenvalues, fermi_energy), - **self._read_occupations(), - **self._read_projections(selection), - } - - def to_database(self, selection=None, fermi_energy=None) -> Band_DB: - occupations = self._read_occupations() - num_total_occupied = occupations.get("occupations", None) - num_checked_bands = None - if num_total_occupied is not None: - num_checked_bands = np.shape(num_total_occupied)[-1] - num_total_occupied = int( - np.max(np.sum(num_total_occupied > _OCCUPATION_CUTOFF, axis=-1)) - ) - num_occupied_up = occupations.get("occupations_up", None) - num_occupied_down = occupations.get("occupations_down", None) - if num_occupied_up is not None: - num_checked_bands = np.shape(num_occupied_up)[-1] - num_occupied_up = int( - np.max(np.sum(num_occupied_up > _OCCUPATION_CUTOFF, axis=-1)) - ) - if num_occupied_down is not None: - num_occupied_down = int( - np.max(np.sum(num_occupied_down > _OCCUPATION_CUTOFF, axis=-1)) - ) - - raw_fermi_energy = ( - self._raw_band.fermi_energy - if not check.is_none(self._raw_band.fermi_energy) - else None - ) - - return Band_DB( - num_considered_bands=num_checked_bands, - num_occupied_bands=num_total_occupied, - num_occupied_bands_up=num_occupied_up, - num_occupied_bands_down=num_occupied_down, - fermi_energy_raw=raw_fermi_energy, - fermi_energy=fermi_energy or raw_fermi_energy, - ) - - def to_graph(self, selection=None, fermi_energy=None, width=0.5) -> graph.Graph: - projections = self._projections(selection, width) - result = self._dispersion().plot(projections) - result = self._shift_series_by_fermi_energy(result, fermi_energy) - result.ylabel = "Energy (eV)" - return result - - def to_frame(self, selection=None, fermi_energy=None): - return pd.DataFrame(self._extract_relevant_data(selection, fermi_energy)) - - def to_quiver(self, selection, normal=None, supercell=None) -> graph.Graph: - reciprocal_lattice_vectors = self._kpoint()._reciprocal_lattice_vectors() - nkp1, nkp2, cut = self._kmesh() - plot_plane = slicing.plane( - reciprocal_lattice_vectors, cut, normal, axis_labels=("b1", "b2", "b3") - ) - options = {"lattice": plot_plane} - if supercell is not None: - options["supercell"] = np.ones(2, dtype=np.int_) * supercell - selector = self._make_selector(self._raw_band.projections) - tree = select.Tree.from_selection(selection) - quiver_plots = [ - graph.Contour( - **self._quiver_plot(selector, sel, nkp1, nkp2, plot_plane), **options - ) - for sel in tree.selections() - ] - return graph.Graph(quiver_plots, title="Spin Texture") - - def selections(self) -> dict: - return self._projector().selections() - - def _is_collinear(self): - return len(self._raw_band.dispersion.eigenvalues) == 2 - - def _is_noncollinear(self): - assert not check.is_none(self._raw_band.projections) - return len(self._raw_band.projections) == 4 - - def _dispersion(self): - return DispersionHandler.from_data(self._raw_band.dispersion) - - def _projector(self): - return ProjectorHandler.from_data(self._raw_band.projectors) - - def _kpoint(self): - return KpointHandler.from_data(self._raw_band.dispersion.kpoints) - - def _projections(self, selection, width): - if selection is None: - return None - check.raise_error_if_not_number( - width, "Width of fat band structure must be a number." - ) - projections = self._read_projections(selection) - spin_projections = projections.get(projector.SPIN_PROJECTION, []) - for label, weight in projections.items(): - if label == projector.SPIN_PROJECTION or label in spin_projections: - continue - weight *= width - return projections - - def _read_projections(self, selection): - return self._projector().project(selection, self._raw_band.projections) - - def _read_occupations(self): - if self._is_collinear(): - return { - "occupations_up": self._raw_band.occupations[0], - "occupations_down": self._raw_band.occupations[1], - } - else: - return {"occupations": self._raw_band.occupations[0]} - - def _shift_dispersion_by_fermi_energy(self, eigenvalues, fermi_energy): - shifted = self._shift_array_by_fermi_energy(eigenvalues, fermi_energy) - if len(shifted) == 2: - return {"bands_up": shifted[0], "bands_down": shifted[1]} - else: - return {"bands": shifted[0]} - - def _shift_series_by_fermi_energy(self, g, fermi_energy): - for series in g.series: - series.y = self._shift_array_by_fermi_energy(series.y, fermi_energy) - return g - - def _shift_array_by_fermi_energy(self, array, fermi_energy): - if fermi_energy is None: - fermi_energy = self._raw_band.fermi_energy - return array - fermi_energy - - def _extract_relevant_data(self, selection, fermi_energy): - need_to_be_repeated = ("kpoint_distances", "kpoint_labels") - relevant_keys = ( - "bands", - "bands_up", - "bands_down", - "occupations", - "occupations_up", - "occupations_down", - ) - data = {} - for key, value in self.to_dict(selection, fermi_energy).items(): - if key in need_to_be_repeated: - data[key] = np.repeat(value, self._raw_band.occupations[0].shape[-1]) - if key in relevant_keys: - data[key] = _to_series(value) - for key, value in self._read_projections(selection).items(): - if key == projector.SPIN_PROJECTION: - continue - data[key] = _to_series(value) - return data - - def _kmesh(self): - try: - nkpx = self._raw_band.dispersion.kpoints.number_x - nkpy = self._raw_band.dispersion.kpoints.number_y - nkpz = self._raw_band.dispersion.kpoints.number_z - if nkpx == 1: - return (nkpy, nkpz, "a") - elif nkpy == 1: - return (nkpx, nkpz, "b") - elif nkpz == 1: - return (nkpx, nkpy, "c") - else: - raise exception.DataMismatch( - f"For spin texture visualisation, the plane normal (a,b,c) to the desired cutting plane must have exactly 1 k-point, but the k-point mesh is {nkpx},{nkpy},{nkpz}. Please adjust the KPOINTS file and re-run VASP." - ) - except exception.NoData: - raise exception.DataMismatch( - "For spin texture visualisation, a k-point grid is assumed, but could not be found for this VASP run." - ) - - def _quiver_plot(self, selector, selection, nkp1, nkp2, plot_plane): - data = selector[selection] - axes = _spin_axes_from_selection(selection) - # Embed the two selected spin components into a 3D Cartesian vector so - # that _project_vectors_to_plane can rotate them into the plot frame. - embedded = np.zeros((3, data.shape[1]), dtype=data.dtype) - embedded[list(axes)] = data - # VASP stores k-points with kx as the fastest-varying (innermost) index. - # Reshape in VASP storage order (slow, fast) then transpose to align with - # the (grid_a, grid_b) layout expected by the plot pipeline. - embedded = embedded.reshape(3, nkp2, nkp1).transpose(0, 2, 1) - projected = slicing._project_vectors_to_plane(plot_plane, embedded) - return {"data": projected, "label": selector.label(selection)} - - def _make_selector(self, projections): - maps = self._projector().to_dict() - maps = { - 1: maps["atom"], - 2: maps["orbital"], - 0: self._spin_map(maps["spin"]), - 4: self._band_map(projections.shape[-1]), - } - return index.Selector( - maps, projections, reduction=_ToQuiverReduction, use_number_labels=True - ) - - def _spin_map(self, spin_map): - if "sigma_x" not in spin_map: - raise exception.DataMismatch( - "System is not noncollinear which is required to visualize spin texture." - ) - return { - "sigma_x~sigma_y": slice(1, 3), - "sigma_x~sigma_z": slice(1, 4, 2), - "sigma_y~sigma_z": slice(2, 4), - "x~y": slice(1, 3), - "x~z": slice(1, 4, 2), - "y~z": slice(2, 4), - "sigma_1~sigma_2": slice(1, 3), - "sigma_1~sigma_3": slice(1, 4, 2), - "sigma_2~sigma_3": slice(2, 4), - } - - def _band_map(self, num_bands): - return {"band": {i + 1: i for i in range(num_bands)}} - - -@quantity("band") -class Band(graph.Mixin): - """The band structure contains the **k** point resolved eigenvalues. - - The most common use case of this class is to produce the electronic band - structure along a path in the Brillouin zone used in a non self consistent - VASP calculation. In some cases you may want to use the `to_dict` function - just to obtain the eigenvalue and projection data though in that case the - **k**-point distances that are calculated are meaningless. - - Examples - -------- - First, we create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - To produce band structure plot use, please check the `to_graph` function for - a more detailed documentation. - - >>> calculation.band.plot() - Graph(series=[Series(x=array(...), y=array(...), label='bands', ...)], - ..., xticks={...}, ..., ylabel='Energy (eV)', ...) - - For your own postprocessing, you can read the band data into a Python dictionary - - >>> calculation.band.read() - {'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...)} - - These methods take additional selections, if you used VASP with :tag:`LORBIT`. - You can inspect possible choices with - - >>> calculation.band.selections() - {'band': ['default', 'kpoints_opt', 'kpoints_wan'], - 'atom': [...], 'orbital': [...], 'spin': [...]} - """ - - def __init__(self, source, quantity_name="band"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_band): - return cls(source=DataSource(raw_band)) - - def _handler_factory(self, raw): - return BandHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - @documentation.format(selection_doc=projector.selection_doc) - def read(self, selection=None, fermi_energy=None) -> dict: - """Read the data into a dictionary. - - You may use this data for your own postprocessing tools. Sometimes you may - want to choose different representations of the electronic band structure or - you want to use the electronic eigenvalues and occupations to compute integrals - over the Brillouin zone. - - We create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> collinear_calculation = demo.calculation(path, selection="collinear") - >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") - - Parameters - ---------- - {selection_doc} - fermi_energy : float - Overwrite the Fermi energy of the band structure calculation with a more - accurate one from a different calculation. This is recommended for metallic - systems where the Fermi energy may be significantly different. - - Returns - ------- - dict - Contains the **k**-point path for plotting band structures with the - eigenvalues shifted to bring the Fermi energy to 0. If available - and a selection is passed, the projections of these bands on the - selected projectors are included. If you specified '''k'''-point labels - in the KPOINTS file, these are returned as well. - - Examples - -------- - Return the **k** points, the electronic eigenvalues, and the Fermi energy as - a Python dictionary - - >>> calculation.band.read() - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...)}} - - Select the p orbitals of the first atom in the POSCAR file: - - >>> calculation.band.read(selection="1(p)") - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...), 'Sr_1_p': array(...)}} - - Select the d orbitals of Sr and Ti: - - >>> calculation.band.read("d(Sr, Ti)") - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...), 'Sr_d': array(...), 'Ti_d': array(...)}} - - For collinear calculations, the spin channels are treated separately - - >>> collinear_calculation.band.read() - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands_up': array(...), - 'bands_down': array(...), 'occupations_up': array(...), - 'occupations_down': array(...)}} - - You can also select particular spin channels, for example the spin-up contribution - of the first three atoms combined - - >>> collinear_calculation.band.read("up(1:3)") - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands_up': array(...), - 'bands_down': array(...), 'occupations_up': array(...), - 'occupations_down': array(...), '1:3_up': array(...)}} - - For noncollinear calculations, the resulting dictionary has the same structure - as for the nonpolarized case - - >>> noncollinear_calculation.band.read() - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...)}} - - If you want to investigate the spin projection of the bands, you can select - particular spin components. Here, we select the x and z components of the spin - - >>> noncollinear_calculation.band.read("sigma_x, sigma_z") - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...), 'sigma_x': array(...), 'sigma_z': array(...), - 'is_spin_projection': ['sigma_x', 'sigma_z']}} - - Add the contribution of three d orbitals - - >>> calculation.band.read("dxy + dxz + dyz") - {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), - 'occupations': array(...), 'dxy + dxz + dyz': array(...)}} - - Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT - file - - >>> calculation.band.read("kpoints_opt") # doctest: +SKIP - {{'kpoint_distances': array(...), 'kpoint_labels': ..., 'fermi_energy': ..., - 'bands': array(...), 'occupations': array(...)}} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandHandler.to_dict, - fermi_energy=fermi_energy, - ) - - def to_dict(self, selection=None, fermi_energy=None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection, fermi_energy=fermi_energy) - - @documentation.format(selection_doc=projector.selection_doc) - def to_graph(self, selection=None, fermi_energy=None, width=0.5) -> graph.Graph: - """Read the data and generate a graph. - - On the x axis, we show the **k** points as distances from the previous ones. - This representation makes sense, if you selected a line mode in the KPOINTS - file. When you provide labels for the **k** points those will be added in the - plot. We show all bands included in the calculation :tag:`NBANDS`. - - If you used the code with :tag:`LORBIT`, you can also plot the projected band - structure. Here, each band will have a linewidth proportional to the projection - of the band on reference orbitals. The maximum width is adjustable with an - argument. - - We create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> collinear_calculation = demo.calculation(path, selection="collinear") - >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") - - Parameters - ---------- - {selection_doc} - fermi_energy : float - Overwrite the Fermi energy of the band structure calculation with a more - accurate one from a different calculation. This is recommended for metallic - systems where the Fermi energy may be significantly different. - width : float - Specifies the width (in eV) of the fatbands if a selection of projections is - specified. If the projection amounts to 100%, the line will be drawn with - this width. - - Returns - ------- - Graph - Figure containing the spin-up and spin-down bands. If a selection - is provided the width of the bands represents the projections of the - bands onto the specified projectors. - - Examples - -------- - Plot the band structure with possible **k** point labels if they have been - provided in the KPOINTS file - - >>> calculation.band.to_graph() - Graph(series=[Series(x=array(...), y=array(...), label='bands', ...)], - ..., xticks={{...}}, ..., ylabel='Energy (eV)', ...) - - Select the p orbitals of the first atom in the POSCAR file: - - >>> calculation.band.to_graph(selection="1(p)") - Graph(series=[Series(..., label='Sr_1_p', weight=array(...), ...)], ...) - - Select the d orbitals of Sr and Ti: - - >>> calculation.band.to_graph("d(Sr, Ti)") - Graph(series=[Series(..., label='Sr_d', ...), Series(..., label='Ti_d', ...)], ...) - - For collinear calculations, the spin channels are treated separately - - >>> collinear_calculation.band.to_graph() - Graph(series=[Series(..., label='up', ...), Series(..., label='down', ...)], ...) - - You can also select particular spin channels, for example the spin-up contribution - of the first three atoms combined - - >>> collinear_calculation.band.to_graph("up(1:3)") - Graph(series=[Series(..., label='1:3_up', ...)], ...) - - For noncollinear calculations, the resulting dictionary has the same structure - as for the nonpolarized case - - >>> noncollinear_calculation.band.to_graph() - Graph(series=[Series(..., label='bands', ...)], ...) - - If you want to investigate the spin projection of the bands, you can select - particular spin components. Here, we select the x component of the spin - - >>> noncollinear_calculation.band.to_graph("sigma_x") - Graph(series=[Series(..., label='sigma_x', ...)], ...) - - Add the contribution of three d orbitals - - >>> calculation.band.to_graph("dxy + dxz + dyz") - Graph(series=[Series(..., label='dxy + dxz + dyz', ...)], ...) - - Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT - file - - >>> calculation.band.to_graph("kpoints_opt") # doctest: +SKIP - Graph(series=[Series(..., label='bands', ...)], ...) - - If you use projections, you can also adjust the width of the lines. Passing - the argument `width=1.0` increases the maximum linewidth to 1 eV - - >>> calculation.band.to_graph("d", width=1.0) - Graph(series=[Series(..., label='d', weight=array(...), ...)], ...) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandHandler.to_graph, - fermi_energy=fermi_energy, - width=width, - ) - - @documentation.format(selection_doc=projector.selection_doc) - def to_frame(self, selection=None, fermi_energy=None): - """Read the data into a DataFrame. - - We create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> collinear_calculation = demo.calculation(path, selection="collinear") - >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") - - Parameters - ---------- - {selection_doc} - fermi_energy : float - Overwrite the Fermi energy of the band structure calculation with a more - accurate one from a different calculation. This is recommended for metallic - systems where the Fermi energy may be significantly different. - - Returns - ------- - pd.DataFrame - Contains the eigenvalues and corresponding occupations for all k-points and - bands. If a selection string is given, in addition the orbital projections - on these bands are returned. - - Examples - -------- - Get the band structure of all bands without projections - - >>> calculation.band.to_frame() - kpoint_distances bands occupations - 0 ... - - Select the p orbitals of the first atom in the POSCAR file: - - >>> calculation.band.to_frame(selection="1(p)") - kpoint_distances bands occupations Sr_1_p - 0 ... - - Select the d orbitals of Sr and Ti: - - >>> calculation.band.to_frame("d(Sr, Ti)") - kpoint_distances bands occupations Sr_d Ti_d - 0 ... - - For collinear calculations, the spin channels are treated separately - - >>> collinear_calculation.band.to_frame() - kpoint_distances bands_up bands_down occupations_up occupations_down - 0 ... - - You can also select particular spin channels, for example the spin-up contribution - of the first three atoms combined - - >>> collinear_calculation.band.to_frame("up(1:3)") - kpoint_distances bands_up ... occupations_down 1:3_up - 0 ... - - For noncollinear calculations, the resulting dictionary has the same structure - as for the nonpolarized case - - >>> noncollinear_calculation.band.to_frame() - kpoint_distances bands occupations - 0 ... - - If you want to investigate the spin projection of the bands, you can select - particular spin components. Here, we select the x and z components of the spin - - >>> noncollinear_calculation.band.to_frame("sigma_x, sigma_z") - kpoint_distances bands occupations sigma_x sigma_z - 0 ... - - Add the contribution of three d orbitals - - >>> calculation.band.to_frame("dxy + dxz + dyz") - kpoint_distances bands occupations dxy + dxz + dyz - 0 ... - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandHandler.to_frame, - fermi_energy=fermi_energy, - ) - - def to_quiver(self, selection, normal=None, supercell=None) -> graph.Graph: - """Generate a quiver plot of spin texture. - - Note that plotting the spin texture requires a special setup of the VASP - calculation. You need to run a noncollinear calculation and a suitable - **k**-point mesh. The **k**-point mesh must be a regular two-dimensional grid, - i.e. one of the three reciprocal lattice directions must not be sampled - (1 in that direction). py4vasp will check this condition and raise an error if - it is not fulfilled. - - The spin texture is represented by arrows in a plane in reciprocal space. The - plane is defined by the two reciprocal lattice vectors that are sampled by the - **k**-point mesh. The normal of this plane can be rotated to align with a - Cartesian axis if desired. The arrows represent the spin expectation value - at each **k** point. You can select which components of the spin are shown - and which bands are included. You can also select particular atoms and orbitals. - - Let us generate some example data do that you can follow along. Please define a - variable `path` with the path to a directory that does not exist yet. - Alternatively, you can use your own data if you have run VASP with an - appropriate k-point mesh. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path, selection="spin_texture") - - Parameters - ---------- - selection - A string specifying the components of the spin and the bands to be included - in the plot. This must be provided and py4vasp will raise an error if it is - not in the correct format. The selection string has the following structure: - - ~(band=) - - where is one of "sigma_x", "sigma_y", or "sigma_z", and - is the index of the band to be included (1-based). For the - spin component, you can also use the alias "x", "y", or "z" and "sigma_1", - "sigma_2", or "sigma_3". - - In addition, you can project onto particular atoms and orbitals similar to - the other methods in this class: - - - To specify the **atom**, you can either use its element name (Si, Al, ...) - or its index as given in the input file (1, 2, ...). For the latter - option it is also possible to specify ranges (e.g. 1:4). - - To select a particular **orbital** you can give a string (s, px, dxz, ...) - or select multiple orbitals by their angular momentum (s, p, d, f). - - If you used a different **k**-point mesh choose "kpoints_opt" or "kpoints_wan" - to select them instead of the default mesh specified in the KPOINTS file. - - You separate multiple selections by commas or whitespace and can nest them using - parenthesis, e.g. `Sr(s, p)` or `s(up), p(down)`. The order of the selections - does not matter, but it is case sensitive to distinguish p (angular momentum - l = 1) from P (phosphorus). - - It is possible to add or subtract different components, e.g., a selection of - "Ti(d) - O(p)" would project onto the d orbitals of Ti and the p orbitals of O - and then compute the difference of these two selections. - - If you are unsure about the specific projections that are available, you can use - - >>> calculation.projector.selections() - {'atom': [...], 'orbital': [...], 'spin': [...]} - - to get a list of all available ones. - normal - Set the Cartesian direction "x", "y", or "z" parallel to which the normal of - the plane is rotated. Alteratively, set it to "auto" to rotate to the closest - Cartesian axis. If you set it to None, the normal will not be considered and - the first remaining lattice vector will be aligned with the x axis instead. - supercell - Replicate the contour plot periodically a given number of times. If you - provide two different numbers, the resulting cell will be the two remaining - lattice vectors multiplied by the specific number. - - Returns - ------- - Graph - A quiver plot for the spin texture in the plane spanned by the 2 remaining - lattice vectors. The arrows represent the spin expectation value at each - **k** point. The plot is replicated periodically according to the specified - supercell. - - Examples - -------- - Plot a projection of the spin texture in reciprocal space, summed over all atoms - and orbitals, for the first band and the x and y components. - - >>> calculation.band.to_quiver("x~y(band=1)") - Graph(series=[Contour(data=array([[[...]]]), ..., label='x~y_band=1', ...)], ..., - title='Spin Texture') - - Select the Ba atom, the third band, the x and z spin components, then sum - over all orbitals: - - >>> calculation.band.to_quiver("Ba(sigma_1~sigma_3(band=3))") - Graph(series=[Contour(..., label='Ba_sigma_1~sigma_3_band=3', ...)], ...) - - Select the Pb atom, s orbital, second band and the x and y spin components: - - >>> calculation.band.to_quiver("Pb(s(band=2(sigma_x~sigma_y)))") - Graph(series=[Contour(..., label='Pb_s_sigma_x~sigma_y_band=2', ...)], ...) - - To plot a 3x3 supercell of the reciprocal lattice plane: - - >>> calculation.band.to_quiver(selection="band=1(x~y)", supercell=3) - Graph(series=[Contour(..., supercell=array([3, 3]), ...)], ...) - - You can also provide different replication numbers for the two remaining - lattice vectors: - - >>> calculation.band.to_quiver(selection="band=1(x~y)", supercell=np.array([2, 4])) - Graph(series=[Contour(..., supercell=array([2, 4]), ...)], ...) - - We can also use KPOINTS_OPT to specify a different mesh. We can use the normal - argument to rotate the plane normal to align with the nearest coordinate axis. - - >>> calculation.band.to_quiver(selection="kpoints_opt(band=1(x~y))", normal="auto") # doctest: +SKIP - Graph(series=[Contour(data=array([[[...]]]), ...)], ...) - - The automatic rotation will only work if one of the reciprocal lattice vectors is - sufficiently close to a Cartesian axis. If this is not the case, you can manually - specify the desired axis. For example, to rotate the plane normal to align with - the x coordinate axis: - - >>> calculation.band.to_quiver(selection="band=1(x~y)", normal="x") - Graph(series=[Contour(data=array([[[...]]]), ...)], ...) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandHandler.to_quiver, - normal=normal, - supercell=supercell, - ) - - def selections(self, selection=None) -> dict: - from py4vasp._raw import definition as raw_module - - handler_selections = merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandHandler.selections, - ) - sources = list(raw_module.selections(self._quantity_name)) - return {self._quantity_name: sources, **handler_selections} - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - self._handler_factory, - BandHandler.to_database, - ) - - -def _to_series(array): - return array.T.flatten() - - -def _spin_axes_from_selection(selection): - """Determine which Cartesian axes the spin pair in *selection* corresponds to.""" - for token in selection: - if ( - not isinstance(token, select.Group) - or token.separator != select.pair_separator - ): - continue - if len(token.group) != 2: - continue - axes = [ - _SPIN_TOKEN_TO_CARTESIAN_AXIS.get(str(member)) for member in token.group - ] - if all(a is not None for a in axes): - return tuple(sorted(axes)) - raise exception.IncorrectUsage( - "Spin Elements must be chosen, but none are given. Please adapt your " - "`selection` argument to include, e.g., `x~y`. You can combine arguments " - "by `arg1(arg2(arg3(...)))`." - ) - - -class _ToQuiverReduction(index.Reduction): - def __init__(self, keys: List): - if not (keys[0]): - raise exception.IncorrectUsage( - "Spin Elements must be chosen, but none are given. Please adapt your `selection` argument to include, e.g., `x~y`. You can combine arguments by `arg1(arg2(arg3(...)))`." - ) - if not (keys[4]): - raise exception.IncorrectUsage( - "A band must be chosen, but none are given. Please adapt your `selection` argument to include, e.g., `band[1]`. You can combine arguments by `arg1(arg2(arg3(...)))`." - ) - super().__init__(keys) - - def __call__(self, array: ArrayLike, axis: Iterable): - axis = tuple(filter(None, axis)) - return np.sum(array, axis=axis) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from __future__ import annotations + +from typing import Any, Iterable, List + +import numpy as np +from numpy.typing import ArrayLike + +from py4vasp import exception +from py4vasp._calculation import projector +from py4vasp._calculation._dispersion import DispersionHandler +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.kpoint import KpointHandler +from py4vasp._calculation.projector import ProjectorHandler +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Band_DB +from py4vasp._third_party import graph +from py4vasp._util import ( + check, + documentation, + import_, + index, + select, + slicing, +) + +pd = import_.optional("pandas") +pretty = import_.optional("IPython.lib.pretty") + +_OCCUPATION_CUTOFF = 1e-2 + +# Cartesian axis (x=0, y=1, z=2) corresponding to each accepted spin-component token. +_SPIN_TOKEN_TO_CARTESIAN_AXIS = { + "x": 0, + "sigma_x": 0, + "sigma_1": 0, + "y": 1, + "sigma_y": 1, + "sigma_2": 1, + "z": 2, + "sigma_z": 2, + "sigma_3": 2, +} + + +class BandHandler: + """Handler for electronic band structure data.""" + + def __init__(self, raw_band: raw.Band): + self._raw_band = raw_band + + @classmethod + def from_data(cls, raw_band: raw.Band) -> "BandHandler": + return cls(raw_band) + + def __str__(self) -> str: + return f""" +{"spin polarized" if self._is_collinear() else ""} band data: + {self._raw_band.dispersion.eigenvalues.shape[1]} k-points + {self._raw_band.dispersion.eigenvalues.shape[2]} bands +{str(self._projector())} + """.strip() + + def to_dict(self, selection=None, fermi_energy=None) -> dict[str, Any]: + dispersion = self._dispersion().to_dict() + eigenvalues = dispersion.pop("eigenvalues") + return { + **dispersion, + "fermi_energy": self._raw_band.fermi_energy, + **self._shift_dispersion_by_fermi_energy(eigenvalues, fermi_energy), + **self._read_occupations(), + **self._read_projections(selection), + } + + def to_database(self, selection=None, fermi_energy=None) -> Band_DB: + occupations = self._read_occupations() + num_total_occupied = occupations.get("occupations", None) + num_checked_bands = None + if num_total_occupied is not None: + num_checked_bands = np.shape(num_total_occupied)[-1] + num_total_occupied = int( + np.max(np.sum(num_total_occupied > _OCCUPATION_CUTOFF, axis=-1)) + ) + num_occupied_up = occupations.get("occupations_up", None) + num_occupied_down = occupations.get("occupations_down", None) + if num_occupied_up is not None: + num_checked_bands = np.shape(num_occupied_up)[-1] + num_occupied_up = int( + np.max(np.sum(num_occupied_up > _OCCUPATION_CUTOFF, axis=-1)) + ) + if num_occupied_down is not None: + num_occupied_down = int( + np.max(np.sum(num_occupied_down > _OCCUPATION_CUTOFF, axis=-1)) + ) + + raw_fermi_energy = ( + self._raw_band.fermi_energy + if not check.is_none(self._raw_band.fermi_energy) + else None + ) + + return Band_DB( + num_considered_bands=num_checked_bands, + num_occupied_bands=num_total_occupied, + num_occupied_bands_up=num_occupied_up, + num_occupied_bands_down=num_occupied_down, + fermi_energy_raw=raw_fermi_energy, + fermi_energy=fermi_energy or raw_fermi_energy, + ) + + def to_graph(self, selection=None, fermi_energy=None, width=0.5) -> graph.Graph: + projections = self._projections(selection, width) + result = self._dispersion().plot(projections) + result = self._shift_series_by_fermi_energy(result, fermi_energy) + result.ylabel = "Energy (eV)" + return result + + def to_frame(self, selection=None, fermi_energy=None): + return pd.DataFrame(self._extract_relevant_data(selection, fermi_energy)) + + def to_quiver(self, selection, normal=None, supercell=None) -> graph.Graph: + reciprocal_lattice_vectors = self._kpoint()._reciprocal_lattice_vectors() + nkp1, nkp2, cut = self._kmesh() + plot_plane = slicing.plane( + reciprocal_lattice_vectors, cut, normal, axis_labels=("b1", "b2", "b3") + ) + options = {"lattice": plot_plane} + if supercell is not None: + options["supercell"] = np.ones(2, dtype=np.int_) * supercell + selector = self._make_selector(self._raw_band.projections) + tree = select.Tree.from_selection(selection) + quiver_plots = [ + graph.Contour( + **self._quiver_plot(selector, sel, nkp1, nkp2, plot_plane), **options + ) + for sel in tree.selections() + ] + return graph.Graph(quiver_plots, title="Spin Texture") + + def selections(self) -> dict: + return self._projector().selections() + + def _is_collinear(self): + return len(self._raw_band.dispersion.eigenvalues) == 2 + + def _is_noncollinear(self): + assert not check.is_none(self._raw_band.projections) + return len(self._raw_band.projections) == 4 + + def _dispersion(self): + return DispersionHandler.from_data(self._raw_band.dispersion) + + def _projector(self): + return ProjectorHandler.from_data(self._raw_band.projectors) + + def _kpoint(self): + return KpointHandler.from_data(self._raw_band.dispersion.kpoints) + + def _projections(self, selection, width): + if selection is None: + return None + check.raise_error_if_not_number( + width, "Width of fat band structure must be a number." + ) + projections = self._read_projections(selection) + spin_projections = projections.get(projector.SPIN_PROJECTION, []) + for label, weight in projections.items(): + if label == projector.SPIN_PROJECTION or label in spin_projections: + continue + weight *= width + return projections + + def _read_projections(self, selection): + return self._projector().project(selection, self._raw_band.projections) + + def _read_occupations(self): + if self._is_collinear(): + return { + "occupations_up": self._raw_band.occupations[0], + "occupations_down": self._raw_band.occupations[1], + } + else: + return {"occupations": self._raw_band.occupations[0]} + + def _shift_dispersion_by_fermi_energy(self, eigenvalues, fermi_energy): + shifted = self._shift_array_by_fermi_energy(eigenvalues, fermi_energy) + if len(shifted) == 2: + return {"bands_up": shifted[0], "bands_down": shifted[1]} + else: + return {"bands": shifted[0]} + + def _shift_series_by_fermi_energy(self, g, fermi_energy): + for series in g.series: + series.y = self._shift_array_by_fermi_energy(series.y, fermi_energy) + return g + + def _shift_array_by_fermi_energy(self, array, fermi_energy): + if fermi_energy is None: + fermi_energy = self._raw_band.fermi_energy + return array - fermi_energy + + def _extract_relevant_data(self, selection, fermi_energy): + need_to_be_repeated = ("kpoint_distances", "kpoint_labels") + relevant_keys = ( + "bands", + "bands_up", + "bands_down", + "occupations", + "occupations_up", + "occupations_down", + ) + data = {} + for key, value in self.to_dict(selection, fermi_energy).items(): + if key in need_to_be_repeated: + data[key] = np.repeat(value, self._raw_band.occupations[0].shape[-1]) + if key in relevant_keys: + data[key] = _to_series(value) + for key, value in self._read_projections(selection).items(): + if key == projector.SPIN_PROJECTION: + continue + data[key] = _to_series(value) + return data + + def _kmesh(self): + try: + nkpx = self._raw_band.dispersion.kpoints.number_x + nkpy = self._raw_band.dispersion.kpoints.number_y + nkpz = self._raw_band.dispersion.kpoints.number_z + if nkpx == 1: + return (nkpy, nkpz, "a") + elif nkpy == 1: + return (nkpx, nkpz, "b") + elif nkpz == 1: + return (nkpx, nkpy, "c") + else: + raise exception.DataMismatch( + f"For spin texture visualisation, the plane normal (a,b,c) to the desired cutting plane must have exactly 1 k-point, but the k-point mesh is {nkpx},{nkpy},{nkpz}. Please adjust the KPOINTS file and re-run VASP." + ) + except exception.NoData: + raise exception.DataMismatch( + "For spin texture visualisation, a k-point grid is assumed, but could not be found for this VASP run." + ) + + def _quiver_plot(self, selector, selection, nkp1, nkp2, plot_plane): + data = selector[selection] + axes = _spin_axes_from_selection(selection) + # Embed the two selected spin components into a 3D Cartesian vector so + # that _project_vectors_to_plane can rotate them into the plot frame. + embedded = np.zeros((3, data.shape[1]), dtype=data.dtype) + embedded[list(axes)] = data + # VASP stores k-points with kx as the fastest-varying (innermost) index. + # Reshape in VASP storage order (slow, fast) then transpose to align with + # the (grid_a, grid_b) layout expected by the plot pipeline. + embedded = embedded.reshape(3, nkp2, nkp1).transpose(0, 2, 1) + projected = slicing._project_vectors_to_plane(plot_plane, embedded) + return {"data": projected, "label": selector.label(selection)} + + def _make_selector(self, projections): + maps = self._projector().to_dict() + maps = { + 1: maps["atom"], + 2: maps["orbital"], + 0: self._spin_map(maps["spin"]), + 4: self._band_map(projections.shape[-1]), + } + return index.Selector( + maps, projections, reduction=_ToQuiverReduction, use_number_labels=True + ) + + def _spin_map(self, spin_map): + if "sigma_x" not in spin_map: + raise exception.DataMismatch( + "System is not noncollinear which is required to visualize spin texture." + ) + return { + "sigma_x~sigma_y": slice(1, 3), + "sigma_x~sigma_z": slice(1, 4, 2), + "sigma_y~sigma_z": slice(2, 4), + "x~y": slice(1, 3), + "x~z": slice(1, 4, 2), + "y~z": slice(2, 4), + "sigma_1~sigma_2": slice(1, 3), + "sigma_1~sigma_3": slice(1, 4, 2), + "sigma_2~sigma_3": slice(2, 4), + } + + def _band_map(self, num_bands): + return {"band": {i + 1: i for i in range(num_bands)}} + + +@quantity("band") +class Band(graph.Mixin): + """The band structure contains the **k** point resolved eigenvalues. + + The most common use case of this class is to produce the electronic band + structure along a path in the Brillouin zone used in a non self consistent + VASP calculation. In some cases you may want to use the `to_dict` function + just to obtain the eigenvalue and projection data though in that case the + **k**-point distances that are calculated are meaningless. + + Examples + -------- + First, we create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + To produce band structure plot use, please check the `to_graph` function for + a more detailed documentation. + + >>> calculation.band.plot() + Graph(series=[Series(x=array(...), y=array(...), label='bands', ...)], + ..., xticks={...}, ..., ylabel='Energy (eV)', ...) + + For your own postprocessing, you can read the band data into a Python dictionary + + >>> calculation.band.read() + {'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...)} + + These methods take additional selections, if you used VASP with :tag:`LORBIT`. + You can inspect possible choices with + + >>> calculation.band.selections() + {'band': ['default', 'kpoints_opt', 'kpoints_wan'], + 'atom': [...], 'orbital': [...], 'spin': [...]} + """ + + def __init__(self, source, quantity_name="band"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_band): + return cls(source=DataSource(raw_band)) + + def _handler_factory(self, raw): + return BandHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + @documentation.format(selection_doc=projector.selection_doc) + def read(self, selection=None, fermi_energy=None) -> dict: + """Read the data into a dictionary. + + You may use this data for your own postprocessing tools. Sometimes you may + want to choose different representations of the electronic band structure or + you want to use the electronic eigenvalues and occupations to compute integrals + over the Brillouin zone. + + We create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> collinear_calculation = demo.calculation(path, selection="collinear") + >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") + + Parameters + ---------- + {selection_doc} + fermi_energy : float + Overwrite the Fermi energy of the band structure calculation with a more + accurate one from a different calculation. This is recommended for metallic + systems where the Fermi energy may be significantly different. + + Returns + ------- + dict + Contains the **k**-point path for plotting band structures with the + eigenvalues shifted to bring the Fermi energy to 0. If available + and a selection is passed, the projections of these bands on the + selected projectors are included. If you specified '''k'''-point labels + in the KPOINTS file, these are returned as well. + + Examples + -------- + Return the **k** points, the electronic eigenvalues, and the Fermi energy as + a Python dictionary + + >>> calculation.band.read() + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...)}} + + Select the p orbitals of the first atom in the POSCAR file: + + >>> calculation.band.read(selection="1(p)") + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...), 'Sr_1_p': array(...)}} + + Select the d orbitals of Sr and Ti: + + >>> calculation.band.read("d(Sr, Ti)") + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...), 'Sr_d': array(...), 'Ti_d': array(...)}} + + For collinear calculations, the spin channels are treated separately + + >>> collinear_calculation.band.read() + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands_up': array(...), + 'bands_down': array(...), 'occupations_up': array(...), + 'occupations_down': array(...)}} + + You can also select particular spin channels, for example the spin-up contribution + of the first three atoms combined + + >>> collinear_calculation.band.read("up(1:3)") + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands_up': array(...), + 'bands_down': array(...), 'occupations_up': array(...), + 'occupations_down': array(...), '1:3_up': array(...)}} + + For noncollinear calculations, the resulting dictionary has the same structure + as for the nonpolarized case + + >>> noncollinear_calculation.band.read() + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...)}} + + If you want to investigate the spin projection of the bands, you can select + particular spin components. Here, we select the x and z components of the spin + + >>> noncollinear_calculation.band.read("sigma_x, sigma_z") + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...), 'sigma_x': array(...), 'sigma_z': array(...), + 'is_spin_projection': ['sigma_x', 'sigma_z']}} + + Add the contribution of three d orbitals + + >>> calculation.band.read("dxy + dxz + dyz") + {{'kpoint_distances': array(...), 'fermi_energy': ..., 'bands': array(...), + 'occupations': array(...), 'dxy + dxz + dyz': array(...)}} + + Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT + file + + >>> calculation.band.read("kpoints_opt") # doctest: +SKIP + {{'kpoint_distances': array(...), 'kpoint_labels': ..., 'fermi_energy': ..., + 'bands': array(...), 'occupations': array(...)}} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandHandler.to_dict, + fermi_energy=fermi_energy, + ) + + def to_dict(self, selection=None, fermi_energy=None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection, fermi_energy=fermi_energy) + + @documentation.format(selection_doc=projector.selection_doc) + def to_graph(self, selection=None, fermi_energy=None, width=0.5) -> graph.Graph: + """Read the data and generate a graph. + + On the x axis, we show the **k** points as distances from the previous ones. + This representation makes sense, if you selected a line mode in the KPOINTS + file. When you provide labels for the **k** points those will be added in the + plot. We show all bands included in the calculation :tag:`NBANDS`. + + If you used the code with :tag:`LORBIT`, you can also plot the projected band + structure. Here, each band will have a linewidth proportional to the projection + of the band on reference orbitals. The maximum width is adjustable with an + argument. + + We create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> collinear_calculation = demo.calculation(path, selection="collinear") + >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") + + Parameters + ---------- + {selection_doc} + fermi_energy : float + Overwrite the Fermi energy of the band structure calculation with a more + accurate one from a different calculation. This is recommended for metallic + systems where the Fermi energy may be significantly different. + width : float + Specifies the width (in eV) of the fatbands if a selection of projections is + specified. If the projection amounts to 100%, the line will be drawn with + this width. + + Returns + ------- + Graph + Figure containing the spin-up and spin-down bands. If a selection + is provided the width of the bands represents the projections of the + bands onto the specified projectors. + + Examples + -------- + Plot the band structure with possible **k** point labels if they have been + provided in the KPOINTS file + + >>> calculation.band.to_graph() + Graph(series=[Series(x=array(...), y=array(...), label='bands', ...)], + ..., xticks={{...}}, ..., ylabel='Energy (eV)', ...) + + Select the p orbitals of the first atom in the POSCAR file: + + >>> calculation.band.to_graph(selection="1(p)") + Graph(series=[Series(..., label='Sr_1_p', weight=array(...), ...)], ...) + + Select the d orbitals of Sr and Ti: + + >>> calculation.band.to_graph("d(Sr, Ti)") + Graph(series=[Series(..., label='Sr_d', ...), Series(..., label='Ti_d', ...)], ...) + + For collinear calculations, the spin channels are treated separately + + >>> collinear_calculation.band.to_graph() + Graph(series=[Series(..., label='up', ...), Series(..., label='down', ...)], ...) + + You can also select particular spin channels, for example the spin-up contribution + of the first three atoms combined + + >>> collinear_calculation.band.to_graph("up(1:3)") + Graph(series=[Series(..., label='1:3_up', ...)], ...) + + For noncollinear calculations, the resulting dictionary has the same structure + as for the nonpolarized case + + >>> noncollinear_calculation.band.to_graph() + Graph(series=[Series(..., label='bands', ...)], ...) + + If you want to investigate the spin projection of the bands, you can select + particular spin components. Here, we select the x component of the spin + + >>> noncollinear_calculation.band.to_graph("sigma_x") + Graph(series=[Series(..., label='sigma_x', ...)], ...) + + Add the contribution of three d orbitals + + >>> calculation.band.to_graph("dxy + dxz + dyz") + Graph(series=[Series(..., label='dxy + dxz + dyz', ...)], ...) + + Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT + file + + >>> calculation.band.to_graph("kpoints_opt") # doctest: +SKIP + Graph(series=[Series(..., label='bands', ...)], ...) + + If you use projections, you can also adjust the width of the lines. Passing + the argument `width=1.0` increases the maximum linewidth to 1 eV + + >>> calculation.band.to_graph("d", width=1.0) + Graph(series=[Series(..., label='d', weight=array(...), ...)], ...) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandHandler.to_graph, + fermi_energy=fermi_energy, + width=width, + ) + + @documentation.format(selection_doc=projector.selection_doc) + def to_frame(self, selection=None, fermi_energy=None): + """Read the data into a DataFrame. + + We create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> collinear_calculation = demo.calculation(path, selection="collinear") + >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") + + Parameters + ---------- + {selection_doc} + fermi_energy : float + Overwrite the Fermi energy of the band structure calculation with a more + accurate one from a different calculation. This is recommended for metallic + systems where the Fermi energy may be significantly different. + + Returns + ------- + pd.DataFrame + Contains the eigenvalues and corresponding occupations for all k-points and + bands. If a selection string is given, in addition the orbital projections + on these bands are returned. + + Examples + -------- + Get the band structure of all bands without projections + + >>> calculation.band.to_frame() + kpoint_distances bands occupations + 0 ... + + Select the p orbitals of the first atom in the POSCAR file: + + >>> calculation.band.to_frame(selection="1(p)") + kpoint_distances bands occupations Sr_1_p + 0 ... + + Select the d orbitals of Sr and Ti: + + >>> calculation.band.to_frame("d(Sr, Ti)") + kpoint_distances bands occupations Sr_d Ti_d + 0 ... + + For collinear calculations, the spin channels are treated separately + + >>> collinear_calculation.band.to_frame() + kpoint_distances bands_up bands_down occupations_up occupations_down + 0 ... + + You can also select particular spin channels, for example the spin-up contribution + of the first three atoms combined + + >>> collinear_calculation.band.to_frame("up(1:3)") + kpoint_distances bands_up ... occupations_down 1:3_up + 0 ... + + For noncollinear calculations, the resulting dictionary has the same structure + as for the nonpolarized case + + >>> noncollinear_calculation.band.to_frame() + kpoint_distances bands occupations + 0 ... + + If you want to investigate the spin projection of the bands, you can select + particular spin components. Here, we select the x and z components of the spin + + >>> noncollinear_calculation.band.to_frame("sigma_x, sigma_z") + kpoint_distances bands occupations sigma_x sigma_z + 0 ... + + Add the contribution of three d orbitals + + >>> calculation.band.to_frame("dxy + dxz + dyz") + kpoint_distances bands occupations dxy + dxz + dyz + 0 ... + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandHandler.to_frame, + fermi_energy=fermi_energy, + ) + + def to_quiver(self, selection, normal=None, supercell=None) -> graph.Graph: + """Generate a quiver plot of spin texture. + + Note that plotting the spin texture requires a special setup of the VASP + calculation. You need to run a noncollinear calculation and a suitable + **k**-point mesh. The **k**-point mesh must be a regular two-dimensional grid, + i.e. one of the three reciprocal lattice directions must not be sampled + (1 in that direction). py4vasp will check this condition and raise an error if + it is not fulfilled. + + The spin texture is represented by arrows in a plane in reciprocal space. The + plane is defined by the two reciprocal lattice vectors that are sampled by the + **k**-point mesh. The normal of this plane can be rotated to align with a + Cartesian axis if desired. The arrows represent the spin expectation value + at each **k** point. You can select which components of the spin are shown + and which bands are included. You can also select particular atoms and orbitals. + + Let us generate some example data do that you can follow along. Please define a + variable `path` with the path to a directory that does not exist yet. + Alternatively, you can use your own data if you have run VASP with an + appropriate k-point mesh. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path, selection="spin_texture") + + Parameters + ---------- + selection + A string specifying the components of the spin and the bands to be included + in the plot. This must be provided and py4vasp will raise an error if it is + not in the correct format. The selection string has the following structure: + + ~(band=) + + where is one of "sigma_x", "sigma_y", or "sigma_z", and + is the index of the band to be included (1-based). For the + spin component, you can also use the alias "x", "y", or "z" and "sigma_1", + "sigma_2", or "sigma_3". + + In addition, you can project onto particular atoms and orbitals similar to + the other methods in this class: + + - To specify the **atom**, you can either use its element name (Si, Al, ...) + or its index as given in the input file (1, 2, ...). For the latter + option it is also possible to specify ranges (e.g. 1:4). + - To select a particular **orbital** you can give a string (s, px, dxz, ...) + or select multiple orbitals by their angular momentum (s, p, d, f). + - If you used a different **k**-point mesh choose "kpoints_opt" or "kpoints_wan" + to select them instead of the default mesh specified in the KPOINTS file. + + You separate multiple selections by commas or whitespace and can nest them using + parenthesis, e.g. `Sr(s, p)` or `s(up), p(down)`. The order of the selections + does not matter, but it is case sensitive to distinguish p (angular momentum + l = 1) from P (phosphorus). + + It is possible to add or subtract different components, e.g., a selection of + "Ti(d) - O(p)" would project onto the d orbitals of Ti and the p orbitals of O + and then compute the difference of these two selections. + + If you are unsure about the specific projections that are available, you can use + + >>> calculation.projector.selections() + {'atom': [...], 'orbital': [...], 'spin': [...]} + + to get a list of all available ones. + normal + Set the Cartesian direction "x", "y", or "z" parallel to which the normal of + the plane is rotated. Alteratively, set it to "auto" to rotate to the closest + Cartesian axis. If you set it to None, the normal will not be considered and + the first remaining lattice vector will be aligned with the x axis instead. + supercell + Replicate the contour plot periodically a given number of times. If you + provide two different numbers, the resulting cell will be the two remaining + lattice vectors multiplied by the specific number. + + Returns + ------- + Graph + A quiver plot for the spin texture in the plane spanned by the 2 remaining + lattice vectors. The arrows represent the spin expectation value at each + **k** point. The plot is replicated periodically according to the specified + supercell. + + Examples + -------- + Plot a projection of the spin texture in reciprocal space, summed over all atoms + and orbitals, for the first band and the x and y components. + + >>> calculation.band.to_quiver("x~y(band=1)") + Graph(series=[Contour(data=array([[[...]]]), ..., label='x~y_band=1', ...)], ..., + title='Spin Texture') + + Select the Ba atom, the third band, the x and z spin components, then sum + over all orbitals: + + >>> calculation.band.to_quiver("Ba(sigma_1~sigma_3(band=3))") + Graph(series=[Contour(..., label='Ba_sigma_1~sigma_3_band=3', ...)], ...) + + Select the Pb atom, s orbital, second band and the x and y spin components: + + >>> calculation.band.to_quiver("Pb(s(band=2(sigma_x~sigma_y)))") + Graph(series=[Contour(..., label='Pb_s_sigma_x~sigma_y_band=2', ...)], ...) + + To plot a 3x3 supercell of the reciprocal lattice plane: + + >>> calculation.band.to_quiver(selection="band=1(x~y)", supercell=3) + Graph(series=[Contour(..., supercell=array([3, 3]), ...)], ...) + + You can also provide different replication numbers for the two remaining + lattice vectors: + + >>> calculation.band.to_quiver(selection="band=1(x~y)", supercell=np.array([2, 4])) + Graph(series=[Contour(..., supercell=array([2, 4]), ...)], ...) + + We can also use KPOINTS_OPT to specify a different mesh. We can use the normal + argument to rotate the plane normal to align with the nearest coordinate axis. + + >>> calculation.band.to_quiver(selection="kpoints_opt(band=1(x~y))", normal="auto") # doctest: +SKIP + Graph(series=[Contour(data=array([[[...]]]), ...)], ...) + + The automatic rotation will only work if one of the reciprocal lattice vectors is + sufficiently close to a Cartesian axis. If this is not the case, you can manually + specify the desired axis. For example, to rotate the plane normal to align with + the x coordinate axis: + + >>> calculation.band.to_quiver(selection="band=1(x~y)", normal="x") + Graph(series=[Contour(data=array([[[...]]]), ...)], ...) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandHandler.to_quiver, + normal=normal, + supercell=supercell, + ) + + def selections(self, selection=None) -> dict: + from py4vasp._raw import definition as raw_module + + handler_selections = merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandHandler.selections, + ) + sources = list(raw_module.selections(self._quantity_name)) + return {self._quantity_name: sources, **handler_selections} + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + self._handler_factory, + BandHandler.to_database, + ) + + +def _to_series(array): + return array.T.flatten() + + +def _spin_axes_from_selection(selection): + """Determine which Cartesian axes the spin pair in *selection* corresponds to.""" + for token in selection: + if ( + not isinstance(token, select.Group) + or token.separator != select.pair_separator + ): + continue + if len(token.group) != 2: + continue + axes = [ + _SPIN_TOKEN_TO_CARTESIAN_AXIS.get(str(member)) for member in token.group + ] + if all(a is not None for a in axes): + return tuple(sorted(axes)) + raise exception.IncorrectUsage( + "Spin Elements must be chosen, but none are given. Please adapt your " + "`selection` argument to include, e.g., `x~y`. You can combine arguments " + "by `arg1(arg2(arg3(...)))`." + ) + + +class _ToQuiverReduction(index.Reduction): + def __init__(self, keys: List): + if not (keys[0]): + raise exception.IncorrectUsage( + "Spin Elements must be chosen, but none are given. Please adapt your `selection` argument to include, e.g., `x~y`. You can combine arguments by `arg1(arg2(arg3(...)))`." + ) + if not (keys[4]): + raise exception.IncorrectUsage( + "A band must be chosen, but none are given. Please adapt your `selection` argument to include, e.g., `band[1]`. You can combine arguments by `arg1(arg2(arg3(...)))`." + ) + super().__init__(keys) + + def __call__(self, array: ArrayLike, axis: Iterable): + axis = tuple(filter(None, axis)) + return np.sum(array, axis=axis) diff --git a/src/py4vasp/_calculation/bandgap.py b/src/py4vasp/_calculation/bandgap.py index 233e3afc7..d3c536ac4 100644 --- a/src/py4vasp/_calculation/bandgap.py +++ b/src/py4vasp/_calculation/bandgap.py @@ -1,474 +1,472 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy -import itertools -import typing - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, - slice_steps, -) -from py4vasp._raw.data_db import Bandgap_DB -from py4vasp._third_party import graph -from py4vasp._util import convert, documentation, select - - -class _Gap(typing.NamedTuple): - bottom: str - top: str - - -GAPS = { - "fundamental": _Gap("valence band maximum", "conduction band minimum"), - "direct": _Gap("direct gap bottom", "direct gap top"), -} - -COMPONENTS = ("independent", "up", "down") - - -class BandgapHandler: - """Handler for the bandgap quantity. Works with exactly one raw.Bandgap object.""" - - def __init__(self, raw_bandgap: raw.Bandgap, steps=None): - self._raw_bandgap = raw_bandgap - self._steps = steps - - @classmethod - def from_data(cls, raw_bandgap: raw.Bandgap, steps=None) -> "BandgapHandler": - return cls(raw_bandgap, steps=steps) - - def to_dict(self) -> dict: - return { - **self._gap_dict("fundamental"), - **self._kpoint_dict("VBM"), - **self._kpoint_dict("CBM"), - **self._gap_dict("direct"), - **self._kpoint_dict("direct"), - "fermi_energy": self._get("Fermi energy", component=0), - } - - def fundamental(self) -> np.ndarray: - """Return the fundamental bandgap.""" - return self._gap("fundamental", component=0) - - def direct(self) -> np.ndarray: - """Return the direct bandgap.""" - return self._gap("direct", component=0) - - def valence_band_maximum(self) -> np.ndarray: - """Return the valence band maximum.""" - return self._get(GAPS["fundamental"].bottom, component=0) - - def conduction_band_minimum(self) -> np.ndarray: - """Return the conduction band minimum.""" - return self._get(GAPS["fundamental"].top, component=0) - - def to_graph(self, selection="fundamental, direct") -> graph.Graph: - """Plot the direct and fundamental bandgap along the trajectory.""" - series = [self._make_series(*choice) for choice in self._parse(selection)] - return graph.Graph(series, xlabel="Step", ylabel="bandgap (eV)") - - def __str__(self): - template = """\ -Band structure --------------- - {header} -val. band max: {val_band_max} -cond. band min: {cond_band_min} -fundamental gap: {fundamental} -VBM @ kpoint: {kpoint_vbm} -CBM @ kpoint: {kpoint_cbm} - -lower band: {lower_band} -upper band: {upper_band} -direct gap: {direct} -@ kpoint: {kpoint_direct} - -Fermi energy: {fermi_energy}""" - return template.format( - header=self._output_header(), - val_band_max=self._output_energy("valence band maximum"), - cond_band_min=self._output_energy("conduction band minimum"), - fundamental=self._output_gap("fundamental"), - kpoint_vbm=self._output_kpoint("VBM"), - kpoint_cbm=self._output_kpoint("CBM"), - lower_band=self._output_energy("direct gap bottom"), - upper_band=self._output_energy("direct gap top"), - direct=self._output_gap("direct"), - kpoint_direct=self._output_kpoint("direct"), - fermi_energy=self._output_energy("Fermi energy", component=slice(0, 1)), - ) - - def to_database(self) -> dict: - bandgap_dict = { - "valence_band_maximum": self._output_energy( - "valence band maximum", to_string=False - ), - "conduction_band_minimum": self._output_energy( - "conduction band minimum", to_string=False - ), - "fundamental_bandgap": self._output_gap("fundamental", to_string=False), - "kpoint_vbm": self._output_kpoint("VBM", to_string=False), - "kpoint_cbm": self._output_kpoint("CBM", to_string=False), - "lower_band_direct_bandgap": self._output_energy( - "direct gap bottom", to_string=False - ), - "upper_band_direct_bandgap": self._output_energy( - "direct gap top", to_string=False - ), - "direct_bandgap": self._output_gap("direct", to_string=False), - "kpoint_direct_bandgap": self._output_kpoint("direct", to_string=False), - } - - final_dict = {} - for k, v in bandgap_dict.items(): - final_dict[f"{k}_spin_independent"] = ( - v[0] if not isinstance(v, float) else float(v) - ) - final_dict[f"{k}_spin_up"] = ( - (v[1] if not isinstance(v[1], float) else float(v[1])) - if self._spin_polarized() - else None - ) - final_dict[f"{k}_spin_down"] = ( - (v[2] if not isinstance(v[2], float) else float(v[2])) - if self._spin_polarized() - else None - ) - return Bandgap_DB(**final_dict) - - # --- Private helpers --- - - def _output_header(self): - if self._spin_polarized(): - return " spin independent spin component 1 spin component 2" - else: - return " spin independent" - - def _output_energy(self, label, component=slice(None), to_string=True): - energies = self._get(label, steps=self._last_step_in_slice, component=component) - if not (to_string): - return energies - return (9 * " ").join(map("{:20.6f}".format, energies)) - - def _output_gap(self, label, to_string=True): - gaps = self._gap(label, steps=self._last_step_in_slice) - if not (to_string): - return gaps - return (9 * " ").join(map("{:20.6f}".format, gaps)) - - def _output_kpoint(self, label, to_string=True): - kpoints = self._kpoint(label, steps=self._last_step_in_slice) - to_string_convert = lambda kpoint: " ".join(map("{:8.4f}".format, kpoint)) - if not (to_string): - return np.array(kpoints).round(decimals=10).tolist() - return " " + " ".join(map(to_string_convert, kpoints)) - - @property - def _last_step_in_slice(self): - if self._steps is None or self._steps == -1: - return -1 - if isinstance(self._steps, slice): - return (self._steps.stop or 0) - 1 - return self._steps - - @property - def _slice(self): - steps = self._steps - if steps is None or steps == -1: - return slice(-1, None) - if isinstance(steps, slice): - return steps - return slice(steps, steps + 1) - - def _gap_dict(self, label): - gaps = self._gap(label).T - return {f"{label}{suffix}": gap for gap, suffix in zip(gaps, self._suffixes())} - - def _kpoint_dict(self, label): - kpoint = self._kpoint(label) - return { - f"kpoint_{label}{suffix}": kpoint[..., i, :] - for i, suffix in enumerate(self._suffixes()) - } - - def _suffixes(self): - return ("", "_up", "_down") if self._spin_polarized() else ("",) - - def _parse(self, selection): - tree = select.Tree.from_selection(selection) - for selection in tree.selections(): - self._raise_error_if_unused_selection(selection) - components = self._parse_components(selection) - labels = self._parse_labels(selection) - yield from itertools.product(labels, components) - - def _raise_error_if_unused_selection(self, selection): - if rest := set(selection).difference(GAPS).difference(COMPONENTS): - raise exception.IncorrectUsage( - f"A part of your selection {rest} could not be mapped to a valid selection" - ) - - def _parse_components(self, selection): - components = set(COMPONENTS).intersection(selection) - if not components: - components = ("independent",) - elif not self._spin_polarized(): - raise exception.IncorrectUsage( - f"You selected a component {components} but the VASP calculation did not include spin polarization." - ) - return components - - def _parse_labels(self, selection): - labels = set(GAPS).intersection(selection) - if not labels: - labels = GAPS.keys() - elif len(labels) > 1: - raise exception.IncorrectUsage( - f"Two conflicting labels selected {labels}. Please check your input." - ) - return labels - - def _make_series(self, label, component): - steps = np.arange(len(self._raw_bandgap.values))[self._slice] + 1 - gaps = self._gap(label, component=COMPONENTS.index(component)) - if component != "independent": - label = f"{label}_{component}" - return graph.Series(steps, np.atleast_1d(gaps), label) - - def _spin_polarized(self): - return self._raw_bandgap.values.shape[1] == 3 - - def _gap(self, label, **kwargs): - top = self._get(GAPS[label].top, **kwargs) - bottom = self._get(GAPS[label].bottom, **kwargs) - return top - bottom - - def _kpoint(self, label, **kwargs): - kpoint = [ - self._get(f"kx ({label})", **kwargs), - self._get(f"ky ({label})", **kwargs), - self._get(f"kz ({label})", **kwargs), - ] - return np.moveaxis(kpoint, 0, -1) - - def _get(self, desired_label, *, steps=None, component=slice(None)): - steps = steps if steps is not None else self._steps - if steps is None: - steps = -1 - return next( - self._raw_bandgap.values[steps, component, index] - for index, label in enumerate(self._raw_bandgap.labels[:]) - if convert.text_to_string(label) == desired_label - ) - - -@quantity("bandgap") -@documentation.format(examples=slice_.examples("bandgap")) -class Bandgap(graph.Mixin): - """This class describes the band extrema during the relaxation or MD simulation. - - The bandgap represents the energy difference between the highest energy electrons - in the valence band and the lowest energy electrons in the conduction band of a - material. The fundamental gap occurs between the energy states of electrons in the - valence and conduction bands irrespective of the **k** point. In contrast, the - direct gap means that transition from valence to conduction band does not change - the **k** momentum. - - To study bandgap the extrema of the valence and conduction band play an important - role. This class reports the valence band maximum as well as the conduction band - minimum. For collinear calculations (ISPIN = 2) all values are reported separately - for both spins as well as ignoring the spin. This simplifies comparison to - experimental data, where the transitions either conserve the spin or not. - - {examples} - """ - - def __init__(self, source, quantity_name: str = "bandgap", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_bandgap: raw.Bandgap): - """Create a Bandgap dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_bandgap)) - - def __getitem__(self, steps) -> "Bandgap": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw_data): - return BandgapHandler.from_data(raw_data, steps=self._steps) - - @documentation.format(examples=slice_.examples("bandgap", "read")) - def read(self, selection: str | None = None) -> dict: - """Read the bandgap data from a VASP relaxation or MD trajectory. - - Returns - ------- - dict - Contains the fundamental and direct gap as well as the coordinates of the - k points where the relevant points in the band structure are. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - @documentation.format(examples=slice_.examples("bandgap", "fundamental")) - def fundamental(self, selection: str | None = None) -> np.ndarray: - """Return the fundamental bandgap. - - The fundamental bandgap is between the maximum of the valence band and the - minimum of the conduction band. - - Returns - ------- - np.ndarray - The value of the bandgap for all selected steps. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.fundamental, - ) - - @documentation.format(examples=slice_.examples("bandgap", "direct")) - def direct(self, selection: str | None = None) -> np.ndarray: - """Return the direct bandgap. - - The direct bandgap is the minimal distance between a valence and conduction - band at a single k point and for a single spin. - - Returns - ------- - np.ndarray - The value of the bandgap for all selected steps. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.direct, - ) - - @documentation.format(examples=slice_.examples("bandgap", "valence_band_maximum")) - def valence_band_maximum(self, selection: str | None = None) -> np.ndarray: - """Return the valence band maximum. - - Returns - ------- - np.ndarray - The value of the valence band maximum for all selected steps. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.valence_band_maximum, - ) - - @documentation.format( - examples=slice_.examples("bandgap", "conduction_band_minimum") - ) - def conduction_band_minimum(self, selection: str | None = None) -> np.ndarray: - """Return the conduction band minimum. - - Returns - ------- - np.ndarray - The value of the conduction band minimum for all selected steps. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.conduction_band_minimum, - ) - - @documentation.format(examples=slice_.examples("bandgap", "to_graph")) - def to_graph(self, selection="fundamental, direct") -> graph.Graph: - """Plot the direct and fundamental bandgap along the trajectory. - - Parameters - ---------- - selection : str - Select which bandgap to include in the plot. By default the fundamental - and the direct one are included. In spin-polarized calculations, you can - also select up or down to obtain the bandgap without spin flips. - - Returns - ------- - Graph - Figure with the ionic step on the x axis and the value of the bandgap on - the y axis. - - {examples} - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.to_graph, - ) - - def __str__(self, selection: str | None = None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - BandgapHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _spin_polarized(self): - """Convenience for tests that access this on the dispatcher.""" - with self._source.access(self._quantity_name) as raw_data: - return raw_data.values.shape[1] == 3 - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - self._handler_factory, - BandgapHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy +import itertools +import typing + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import Bandgap_DB +from py4vasp._third_party import graph +from py4vasp._util import convert, documentation, select + + +class _Gap(typing.NamedTuple): + bottom: str + top: str + + +GAPS = { + "fundamental": _Gap("valence band maximum", "conduction band minimum"), + "direct": _Gap("direct gap bottom", "direct gap top"), +} + +COMPONENTS = ("independent", "up", "down") + + +class BandgapHandler: + """Handler for the bandgap quantity. Works with exactly one raw.Bandgap object.""" + + def __init__(self, raw_bandgap: raw.Bandgap, steps=None): + self._raw_bandgap = raw_bandgap + self._steps = steps + + @classmethod + def from_data(cls, raw_bandgap: raw.Bandgap, steps=None) -> "BandgapHandler": + return cls(raw_bandgap, steps=steps) + + def to_dict(self) -> dict: + return { + **self._gap_dict("fundamental"), + **self._kpoint_dict("VBM"), + **self._kpoint_dict("CBM"), + **self._gap_dict("direct"), + **self._kpoint_dict("direct"), + "fermi_energy": self._get("Fermi energy", component=0), + } + + def fundamental(self) -> np.ndarray: + """Return the fundamental bandgap.""" + return self._gap("fundamental", component=0) + + def direct(self) -> np.ndarray: + """Return the direct bandgap.""" + return self._gap("direct", component=0) + + def valence_band_maximum(self) -> np.ndarray: + """Return the valence band maximum.""" + return self._get(GAPS["fundamental"].bottom, component=0) + + def conduction_band_minimum(self) -> np.ndarray: + """Return the conduction band minimum.""" + return self._get(GAPS["fundamental"].top, component=0) + + def to_graph(self, selection="fundamental, direct") -> graph.Graph: + """Plot the direct and fundamental bandgap along the trajectory.""" + series = [self._make_series(*choice) for choice in self._parse(selection)] + return graph.Graph(series, xlabel="Step", ylabel="bandgap (eV)") + + def __str__(self): + template = """\ +Band structure +-------------- + {header} +val. band max: {val_band_max} +cond. band min: {cond_band_min} +fundamental gap: {fundamental} +VBM @ kpoint: {kpoint_vbm} +CBM @ kpoint: {kpoint_cbm} + +lower band: {lower_band} +upper band: {upper_band} +direct gap: {direct} +@ kpoint: {kpoint_direct} + +Fermi energy: {fermi_energy}""" + return template.format( + header=self._output_header(), + val_band_max=self._output_energy("valence band maximum"), + cond_band_min=self._output_energy("conduction band minimum"), + fundamental=self._output_gap("fundamental"), + kpoint_vbm=self._output_kpoint("VBM"), + kpoint_cbm=self._output_kpoint("CBM"), + lower_band=self._output_energy("direct gap bottom"), + upper_band=self._output_energy("direct gap top"), + direct=self._output_gap("direct"), + kpoint_direct=self._output_kpoint("direct"), + fermi_energy=self._output_energy("Fermi energy", component=slice(0, 1)), + ) + + def to_database(self) -> dict: + bandgap_dict = { + "valence_band_maximum": self._output_energy( + "valence band maximum", to_string=False + ), + "conduction_band_minimum": self._output_energy( + "conduction band minimum", to_string=False + ), + "fundamental_bandgap": self._output_gap("fundamental", to_string=False), + "kpoint_vbm": self._output_kpoint("VBM", to_string=False), + "kpoint_cbm": self._output_kpoint("CBM", to_string=False), + "lower_band_direct_bandgap": self._output_energy( + "direct gap bottom", to_string=False + ), + "upper_band_direct_bandgap": self._output_energy( + "direct gap top", to_string=False + ), + "direct_bandgap": self._output_gap("direct", to_string=False), + "kpoint_direct_bandgap": self._output_kpoint("direct", to_string=False), + } + + final_dict = {} + for k, v in bandgap_dict.items(): + final_dict[f"{k}_spin_independent"] = ( + v[0] if not isinstance(v, float) else float(v) + ) + final_dict[f"{k}_spin_up"] = ( + (v[1] if not isinstance(v[1], float) else float(v[1])) + if self._spin_polarized() + else None + ) + final_dict[f"{k}_spin_down"] = ( + (v[2] if not isinstance(v[2], float) else float(v[2])) + if self._spin_polarized() + else None + ) + return Bandgap_DB(**final_dict) + + # --- Private helpers --- + + def _output_header(self): + if self._spin_polarized(): + return " spin independent spin component 1 spin component 2" + else: + return " spin independent" + + def _output_energy(self, label, component=slice(None), to_string=True): + energies = self._get(label, steps=self._last_step_in_slice, component=component) + if not (to_string): + return energies + return (9 * " ").join(map("{:20.6f}".format, energies)) + + def _output_gap(self, label, to_string=True): + gaps = self._gap(label, steps=self._last_step_in_slice) + if not (to_string): + return gaps + return (9 * " ").join(map("{:20.6f}".format, gaps)) + + def _output_kpoint(self, label, to_string=True): + kpoints = self._kpoint(label, steps=self._last_step_in_slice) + to_string_convert = lambda kpoint: " ".join(map("{:8.4f}".format, kpoint)) + if not (to_string): + return np.array(kpoints).round(decimals=10).tolist() + return " " + " ".join(map(to_string_convert, kpoints)) + + @property + def _last_step_in_slice(self): + if self._steps is None or self._steps == -1: + return -1 + if isinstance(self._steps, slice): + return (self._steps.stop or 0) - 1 + return self._steps + + @property + def _slice(self): + steps = self._steps + if steps is None or steps == -1: + return slice(-1, None) + if isinstance(steps, slice): + return steps + return slice(steps, steps + 1) + + def _gap_dict(self, label): + gaps = self._gap(label).T + return {f"{label}{suffix}": gap for gap, suffix in zip(gaps, self._suffixes())} + + def _kpoint_dict(self, label): + kpoint = self._kpoint(label) + return { + f"kpoint_{label}{suffix}": kpoint[..., i, :] + for i, suffix in enumerate(self._suffixes()) + } + + def _suffixes(self): + return ("", "_up", "_down") if self._spin_polarized() else ("",) + + def _parse(self, selection): + tree = select.Tree.from_selection(selection) + for selection in tree.selections(): + self._raise_error_if_unused_selection(selection) + components = self._parse_components(selection) + labels = self._parse_labels(selection) + yield from itertools.product(labels, components) + + def _raise_error_if_unused_selection(self, selection): + if rest := set(selection).difference(GAPS).difference(COMPONENTS): + raise exception.IncorrectUsage( + f"A part of your selection {rest} could not be mapped to a valid selection" + ) + + def _parse_components(self, selection): + components = set(COMPONENTS).intersection(selection) + if not components: + components = ("independent",) + elif not self._spin_polarized(): + raise exception.IncorrectUsage( + f"You selected a component {components} but the VASP calculation did not include spin polarization." + ) + return components + + def _parse_labels(self, selection): + labels = set(GAPS).intersection(selection) + if not labels: + labels = GAPS.keys() + elif len(labels) > 1: + raise exception.IncorrectUsage( + f"Two conflicting labels selected {labels}. Please check your input." + ) + return labels + + def _make_series(self, label, component): + steps = np.arange(len(self._raw_bandgap.values))[self._slice] + 1 + gaps = self._gap(label, component=COMPONENTS.index(component)) + if component != "independent": + label = f"{label}_{component}" + return graph.Series(steps, np.atleast_1d(gaps), label) + + def _spin_polarized(self): + return self._raw_bandgap.values.shape[1] == 3 + + def _gap(self, label, **kwargs): + top = self._get(GAPS[label].top, **kwargs) + bottom = self._get(GAPS[label].bottom, **kwargs) + return top - bottom + + def _kpoint(self, label, **kwargs): + kpoint = [ + self._get(f"kx ({label})", **kwargs), + self._get(f"ky ({label})", **kwargs), + self._get(f"kz ({label})", **kwargs), + ] + return np.moveaxis(kpoint, 0, -1) + + def _get(self, desired_label, *, steps=None, component=slice(None)): + steps = steps if steps is not None else self._steps + if steps is None: + steps = -1 + return next( + self._raw_bandgap.values[steps, component, index] + for index, label in enumerate(self._raw_bandgap.labels[:]) + if convert.text_to_string(label) == desired_label + ) + + +@quantity("bandgap") +@documentation.format(examples=slice_.examples("bandgap")) +class Bandgap(graph.Mixin): + """This class describes the band extrema during the relaxation or MD simulation. + + The bandgap represents the energy difference between the highest energy electrons + in the valence band and the lowest energy electrons in the conduction band of a + material. The fundamental gap occurs between the energy states of electrons in the + valence and conduction bands irrespective of the **k** point. In contrast, the + direct gap means that transition from valence to conduction band does not change + the **k** momentum. + + To study bandgap the extrema of the valence and conduction band play an important + role. This class reports the valence band maximum as well as the conduction band + minimum. For collinear calculations (ISPIN = 2) all values are reported separately + for both spins as well as ignoring the spin. This simplifies comparison to + experimental data, where the transitions either conserve the spin or not. + + {examples} + """ + + def __init__(self, source, quantity_name: str = "bandgap", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_bandgap: raw.Bandgap): + """Create a Bandgap dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_bandgap)) + + def __getitem__(self, steps) -> "Bandgap": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw_data): + return BandgapHandler.from_data(raw_data, steps=self._steps) + + @documentation.format(examples=slice_.examples("bandgap", "read")) + def read(self, selection: str | None = None) -> dict: + """Read the bandgap data from a VASP relaxation or MD trajectory. + + Returns + ------- + dict + Contains the fundamental and direct gap as well as the coordinates of the + k points where the relevant points in the band structure are. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + @documentation.format(examples=slice_.examples("bandgap", "fundamental")) + def fundamental(self, selection: str | None = None) -> np.ndarray: + """Return the fundamental bandgap. + + The fundamental bandgap is between the maximum of the valence band and the + minimum of the conduction band. + + Returns + ------- + np.ndarray + The value of the bandgap for all selected steps. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.fundamental, + ) + + @documentation.format(examples=slice_.examples("bandgap", "direct")) + def direct(self, selection: str | None = None) -> np.ndarray: + """Return the direct bandgap. + + The direct bandgap is the minimal distance between a valence and conduction + band at a single k point and for a single spin. + + Returns + ------- + np.ndarray + The value of the bandgap for all selected steps. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.direct, + ) + + @documentation.format(examples=slice_.examples("bandgap", "valence_band_maximum")) + def valence_band_maximum(self, selection: str | None = None) -> np.ndarray: + """Return the valence band maximum. + + Returns + ------- + np.ndarray + The value of the valence band maximum for all selected steps. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.valence_band_maximum, + ) + + @documentation.format( + examples=slice_.examples("bandgap", "conduction_band_minimum") + ) + def conduction_band_minimum(self, selection: str | None = None) -> np.ndarray: + """Return the conduction band minimum. + + Returns + ------- + np.ndarray + The value of the conduction band minimum for all selected steps. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.conduction_band_minimum, + ) + + @documentation.format(examples=slice_.examples("bandgap", "to_graph")) + def to_graph(self, selection="fundamental, direct") -> graph.Graph: + """Plot the direct and fundamental bandgap along the trajectory. + + Parameters + ---------- + selection : str + Select which bandgap to include in the plot. By default the fundamental + and the direct one are included. In spin-polarized calculations, you can + also select up or down to obtain the bandgap without spin flips. + + Returns + ------- + Graph + Figure with the ionic step on the x axis and the value of the bandgap on + the y axis. + + {examples} + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.to_graph, + ) + + def __str__(self, selection: str | None = None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + BandgapHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _spin_polarized(self): + """Convenience for tests that access this on the dispatcher.""" + with self._source.access(self._quantity_name) as raw_data: + return raw_data.values.shape[1] == 3 + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + self._handler_factory, + BandgapHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/born_effective_charge.py b/src/py4vasp/_calculation/born_effective_charge.py index 606857010..0361eecd4 100644 --- a/src/py4vasp/_calculation/born_effective_charge.py +++ b/src/py4vasp/_calculation/born_effective_charge.py @@ -1,168 +1,167 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) - -import numpy as np - -from py4vasp import raw -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw.data_db import BornEffectiveCharge_DB -from py4vasp._util import check - - -class BornEffectiveChargeHandler: - """The Born effective charge tensors couple electric field and atomic displacement.""" - - def __init__(self, raw_born_effective_charge: raw.BornEffectiveCharge): - self._raw_born_effective_charge = raw_born_effective_charge - - @classmethod - def from_data( - cls, raw_born_effective_charge: raw.BornEffectiveCharge - ) -> "BornEffectiveChargeHandler": - return cls(raw_born_effective_charge) - - def __str__(self) -> str: - data = self.to_dict() - result = """ -BORN EFFECTIVE CHARGES (including local field effects) (in |e|, cumulative output) ---------------------------------------------------------------------------------- - """.strip() - generator = zip(data["structure"]["elements"], data["charge_tensors"]) - vec_to_string = lambda vec: " ".join(f"{x:11.5f}" for x in vec) - for ion, (element, charge_tensor) in enumerate(generator): - result += f""" -ion {ion + 1:4d} {element} - 1 {vec_to_string(charge_tensor[0])} - 2 {vec_to_string(charge_tensor[1])} - 3 {vec_to_string(charge_tensor[2])}""" - return result - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Read structure information and Born effective charges into a dictionary.""" - return self.to_dict() - - def to_dict(self) -> dict: - """Read structure information and Born effective charges into a dictionary. - - The structural information is added to inform about which atoms are included - in the array. The Born effective charges array contains the mixed second - derivative with respect to an electric field and an atomic displacement for - all atoms and possible directions. - - Returns - ------- - dict - Contains structural information as well as the Born effective charges. - """ - structure = StructureHandler.from_data( - self._raw_born_effective_charge.structure - ) - return { - "structure": structure.to_dict(), - "charge_tensors": self._raw_born_effective_charge.charge_tensors[:], - } - - def to_database(self) -> dict: - """Return Born effective charge data ready for database storage.""" - eigenvalue_max = None - eigenvalue_max_index = None - eigenvalue_min = None - eigenvalue_min_index = None - - if not check.is_none(self._raw_born_effective_charge.charge_tensors): - charge_tensors = self._raw_born_effective_charge.charge_tensors[:] - traces = ( - charge_tensors[:, 0, 0] - + charge_tensors[:, 1, 1] - + charge_tensors[:, 2, 2] - ) - eigenvalue_max = float(np.max(traces)) - eigenvalue_min = float(np.min(traces)) - eigenvalue_max_index = int(np.argmax(traces)) - eigenvalue_min_index = int(np.argmin(traces)) - - return BornEffectiveCharge_DB( - eigenvalue_min=eigenvalue_min, - eigenvalue_min_index=eigenvalue_min_index, - eigenvalue_max=eigenvalue_max, - eigenvalue_max_index=eigenvalue_max_index, - ) - - -@quantity("born_effective_charge") -class BornEffectiveCharge: - """The Born effective charge tensors couple electric field and atomic displacement. - - You can use this class to extract the Born effective charges of a linear - response calculation. The Born effective charges describes the effective charge of - an ion in a crystal lattice when subjected to an external electric field. - These charges account for the displacement of the ion positions in response to the - field, reflecting the distortion of the crystal structure. Born effective charges - help understanding the material's response to external stimuli, such as - piezoelectric and ferroelectric behavior. - """ - - def __init__(self, source, quantity_name: str = "born_effective_charge"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data( - cls, raw_born_effective_charge: raw.BornEffectiveCharge - ) -> "BornEffectiveCharge": - """Create a BornEffectiveCharge dispatcher from raw data.""" - return cls(source=DataSource(raw_born_effective_charge)) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - BornEffectiveChargeHandler.from_data, - BornEffectiveChargeHandler.__str__, - ) - - def read(self) -> dict: - """Read structure information and Born effective charges into a dictionary. - - The structural information is added to inform about which atoms are included - in the array. The Born effective charges array contains the mixed second - derivative with respect to an electric field and an atomic displacement for - all atoms and possible directions. - - Returns - ------- - dict - Contains structural information as well as the Born effective charges. - """ - return merge_default( - self._source, - self._quantity_name, - None, - BornEffectiveChargeHandler.from_data, - BornEffectiveChargeHandler.read, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - BornEffectiveChargeHandler.from_data, - BornEffectiveChargeHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import numpy as np + +from py4vasp import raw +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw.data_db import BornEffectiveCharge_DB +from py4vasp._util import check + + +class BornEffectiveChargeHandler: + """The Born effective charge tensors couple electric field and atomic displacement.""" + + def __init__(self, raw_born_effective_charge: raw.BornEffectiveCharge): + self._raw_born_effective_charge = raw_born_effective_charge + + @classmethod + def from_data( + cls, raw_born_effective_charge: raw.BornEffectiveCharge + ) -> "BornEffectiveChargeHandler": + return cls(raw_born_effective_charge) + + def __str__(self) -> str: + data = self.to_dict() + result = """ +BORN EFFECTIVE CHARGES (including local field effects) (in |e|, cumulative output) +--------------------------------------------------------------------------------- + """.strip() + generator = zip(data["structure"]["elements"], data["charge_tensors"]) + vec_to_string = lambda vec: " ".join(f"{x:11.5f}" for x in vec) + for ion, (element, charge_tensor) in enumerate(generator): + result += f""" +ion {ion + 1:4d} {element} + 1 {vec_to_string(charge_tensor[0])} + 2 {vec_to_string(charge_tensor[1])} + 3 {vec_to_string(charge_tensor[2])}""" + return result + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Read structure information and Born effective charges into a dictionary.""" + return self.to_dict() + + def to_dict(self) -> dict: + """Read structure information and Born effective charges into a dictionary. + + The structural information is added to inform about which atoms are included + in the array. The Born effective charges array contains the mixed second + derivative with respect to an electric field and an atomic displacement for + all atoms and possible directions. + + Returns + ------- + dict + Contains structural information as well as the Born effective charges. + """ + structure = StructureHandler.from_data( + self._raw_born_effective_charge.structure + ) + return { + "structure": structure.to_dict(), + "charge_tensors": self._raw_born_effective_charge.charge_tensors[:], + } + + def to_database(self) -> dict: + """Return Born effective charge data ready for database storage.""" + eigenvalue_max = None + eigenvalue_max_index = None + eigenvalue_min = None + eigenvalue_min_index = None + + if not check.is_none(self._raw_born_effective_charge.charge_tensors): + charge_tensors = self._raw_born_effective_charge.charge_tensors[:] + traces = ( + charge_tensors[:, 0, 0] + + charge_tensors[:, 1, 1] + + charge_tensors[:, 2, 2] + ) + eigenvalue_max = float(np.max(traces)) + eigenvalue_min = float(np.min(traces)) + eigenvalue_max_index = int(np.argmax(traces)) + eigenvalue_min_index = int(np.argmin(traces)) + + return BornEffectiveCharge_DB( + eigenvalue_min=eigenvalue_min, + eigenvalue_min_index=eigenvalue_min_index, + eigenvalue_max=eigenvalue_max, + eigenvalue_max_index=eigenvalue_max_index, + ) + + +@quantity("born_effective_charge") +class BornEffectiveCharge: + """The Born effective charge tensors couple electric field and atomic displacement. + + You can use this class to extract the Born effective charges of a linear + response calculation. The Born effective charges describes the effective charge of + an ion in a crystal lattice when subjected to an external electric field. + These charges account for the displacement of the ion positions in response to the + field, reflecting the distortion of the crystal structure. Born effective charges + help understanding the material's response to external stimuli, such as + piezoelectric and ferroelectric behavior. + """ + + def __init__(self, source, quantity_name: str = "born_effective_charge"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data( + cls, raw_born_effective_charge: raw.BornEffectiveCharge + ) -> "BornEffectiveCharge": + """Create a BornEffectiveCharge dispatcher from raw data.""" + return cls(source=DataSource(raw_born_effective_charge)) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + BornEffectiveChargeHandler.from_data, + BornEffectiveChargeHandler.__str__, + ) + + def read(self) -> dict: + """Read structure information and Born effective charges into a dictionary. + + The structural information is added to inform about which atoms are included + in the array. The Born effective charges array contains the mixed second + derivative with respect to an electric field and an atomic displacement for + all atoms and possible directions. + + Returns + ------- + dict + Contains structural information as well as the Born effective charges. + """ + return merge_default( + self._source, + self._quantity_name, + None, + BornEffectiveChargeHandler.from_data, + BornEffectiveChargeHandler.read, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + BornEffectiveChargeHandler.from_data, + BornEffectiveChargeHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/current_density.py b/src/py4vasp/_calculation/current_density.py index cde55558c..cc009cc0c 100644 --- a/src/py4vasp/_calculation/current_density.py +++ b/src/py4vasp/_calculation/current_density.py @@ -1,328 +1,327 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy -from typing import Optional, Union - -import numpy as np - -from py4vasp import exception -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw import data as raw -from py4vasp._third_party import graph -from py4vasp._util import documentation, import_, slicing -from py4vasp._util.density import SliceArguments, Visualizer - -pretty = import_.optional("IPython.lib.pretty") - -_COMMON_PARAMETERS = f"""selection : str | None = None - Selects which of the possible available currents is used. Check the - `selections` method for all available choices. - -{slicing.PARAMETERS} -supercell : int | np.ndarray | None = None - Replicate the contour plot periodically a given number of times. If you - provide two different numbers, the resulting cell will be the two remaining - lattice vectors multiplied by the specific number.""" - - -class CurrentDensityHandler: - """Handler for current density data — performs all data access and transformation.""" - - def __init__(self, raw_current_density: raw.CurrentDensity): - self._raw_current_density = raw_current_density - - @classmethod - def from_data( - cls, raw_current_density: raw.CurrentDensity - ) -> "CurrentDensityHandler": - return cls(raw_current_density) - - def __str__(self) -> str: - raw_stoichiometry = self._raw_current_density.structure.stoichiometry - stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry) - key = self._raw_current_density.valid_indices[-1] - grid = self._raw_current_density[key].current_density.shape[1:] - return f"""current density: - structure: {pretty.pretty(stoichiometry)} - grid: {grid[2]}, {grid[1]}, {grid[0]} - selections: {", ".join(str(index) for index in self._raw_current_density.valid_indices)}""" - - def to_dict(self) -> dict: - """Read the current density and structural information into a Python dictionary.""" - return { - "structure": self._structure().to_dict(), - **self._read_current_densities(), - } - - def to_database(self) -> dict: - return {} - - def to_contour( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - """Generate a contour plot of current density.""" - label, grid_vector = self._read_current_density(selection) - grid_scalar = np.linalg.norm(grid_vector, axis=-1) - visualizer = Visualizer(self._structure()) - return visualizer.to_contour( - {label: grid_scalar}, - SliceArguments(a, b, c, normal, supercell), - isolevels=False, - ) - - def to_quiver( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - """Generate a quiver plot of current density.""" - label, data = self._read_current_density(selection) - sel_data = np.moveaxis(data, -1, 0) - visualizer = Visualizer(self._structure()) - return visualizer.to_quiver( - {label: sel_data}, SliceArguments(a, b, c, normal, supercell) - ) - - def _structure(self): - return StructureHandler.from_data(self._raw_current_density.structure) - - def _read_current_densities(self): - return dict( - self._read_current_density(key) for key in self._raw_current_density - ) - - def _read_current_density(self, key=None): - key = key or self._raw_current_density.valid_indices[-1] - if key not in self._raw_current_density.valid_indices: - raise exception.IncorrectUsage( - f"Selection {key!r} is not available. " - f"Please use one of {list(self._raw_current_density.valid_indices)}." - ) - return f"current_{key}", self._raw_current_density[key].current_density[:].T - - -@quantity("current_density") -class CurrentDensity: - """Represents current density on the grid in the unit cell. - - A current density j is a vectorial quantity (j_x, j_y, j_z) on every grid point. - It describes how the current flows at every point in space. - - Examples - -------- - - First, we create some example data that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - To produce current density plots, please check the `to_contour` and `to_quiver` - functions for a more detailed documentation. - - To produce a contour plot: - - >>> calculation.current_density.to_contour("nmr", a=0) - Graph(series=[Contour(data=array([[...]]), ..., cut='a', ...)], ...) - - To produce a quiver plot: - - >>> calculation.current_density.to_quiver("nmr", c=0) - Graph(series=[Contour(data=array([[[...]]]), ..., cut='c', ...)], ...) - - For your own postprocessing, you can read the current density data into a Python dict: - - >>> calculation.current_density.read("nmr") - {'structure': {...}, 'current_x': array([[[[...]]]], ...), 'current_y': array([[[[...]]]], ...), 'current_z': array([[[[...]]]], ...)} - - You can inspect possible choices with: - - >>> calculation.current_density.selections("nmr") - {'current_density': ['nmr']} - - Please check the documentation of these methods for more details on how to use them - and which options they provide.""" - - def __init__(self, source, quantity_name="current_density"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_current_density): - return cls(source=DataSource(raw_current_density)) - - def _handler_factory(self, raw): - return CurrentDensityHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - CurrentDensityHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Read the current density and structural information into a Python dictionary. - - Returns - ------- - dict - Contains all available current density data as well as structural information. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - CurrentDensityHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) - def to_contour( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - """Generate a contour plot of current density. - - {plane} - - Parameters - ---------- - {parameters} - - Returns - ------- - Graph - A current density plot in the plane spanned by the 2 remaining lattice - vectors. - - Examples - -------- - - Cut a plane at the origin of the third lattice vector. - - >>> calculation.current_density.to_contour("nmr", c=0) - - Replicate a plane in the middle of the second lattice vector 2 times in each - direction. - - >>> calculation.current_density.to_contour("nmr", b=0.5, supercell=2) - - Take a slice along the first lattice vector and rotate it such that the normal - of the plane aligns with the x axis. - - >>> calculation.current_density.to_contour("nmr", a=0.3, normal="x") - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - CurrentDensityHandler.to_contour, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) - def to_quiver( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - """Generate a quiver plot of current density. - - {plane} - - Parameters - ---------- - {parameters} - - Returns - ------- - Graph - A quiver plot in the plane spanned by the 2 remaining lattice vectors. - - Examples - -------- - - Cut a plane at the origin of the third lattice vector. - - >>> calculation.current_density.to_quiver("nmr", c=0) - - Replicate a plane in the middle of the second lattice vector 2 times in each - direction. - - >>> calculation.current_density.to_quiver("nmr", b=0.5, supercell=2) - - Take a slice along the first lattice vector and rotate it such that the normal - of the plane aligns with the x axis. - - >>> calculation.current_density.to_quiver("nmr", a=0.3, normal="x") - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - CurrentDensityHandler.to_quiver, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - CurrentDensityHandler.from_data, - CurrentDensityHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from typing import Optional, Union + +import numpy as np + +from py4vasp import exception +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw import data as raw +from py4vasp._third_party import graph +from py4vasp._util import documentation, import_, slicing +from py4vasp._util.density import SliceArguments, Visualizer + +pretty = import_.optional("IPython.lib.pretty") + +_COMMON_PARAMETERS = f"""selection : str | None = None + Selects which of the possible available currents is used. Check the + `selections` method for all available choices. + +{slicing.PARAMETERS} +supercell : int | np.ndarray | None = None + Replicate the contour plot periodically a given number of times. If you + provide two different numbers, the resulting cell will be the two remaining + lattice vectors multiplied by the specific number.""" + + +class CurrentDensityHandler: + """Handler for current density data — performs all data access and transformation.""" + + def __init__(self, raw_current_density: raw.CurrentDensity): + self._raw_current_density = raw_current_density + + @classmethod + def from_data( + cls, raw_current_density: raw.CurrentDensity + ) -> "CurrentDensityHandler": + return cls(raw_current_density) + + def __str__(self) -> str: + raw_stoichiometry = self._raw_current_density.structure.stoichiometry + stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry) + key = self._raw_current_density.valid_indices[-1] + grid = self._raw_current_density[key].current_density.shape[1:] + return f"""current density: + structure: {pretty.pretty(stoichiometry)} + grid: {grid[2]}, {grid[1]}, {grid[0]} + selections: {", ".join(str(index) for index in self._raw_current_density.valid_indices)}""" + + def to_dict(self) -> dict: + """Read the current density and structural information into a Python dictionary.""" + return { + "structure": self._structure().to_dict(), + **self._read_current_densities(), + } + + def to_database(self) -> dict: + return {} + + def to_contour( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + """Generate a contour plot of current density.""" + label, grid_vector = self._read_current_density(selection) + grid_scalar = np.linalg.norm(grid_vector, axis=-1) + visualizer = Visualizer(self._structure()) + return visualizer.to_contour( + {label: grid_scalar}, + SliceArguments(a, b, c, normal, supercell), + isolevels=False, + ) + + def to_quiver( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + """Generate a quiver plot of current density.""" + label, data = self._read_current_density(selection) + sel_data = np.moveaxis(data, -1, 0) + visualizer = Visualizer(self._structure()) + return visualizer.to_quiver( + {label: sel_data}, SliceArguments(a, b, c, normal, supercell) + ) + + def _structure(self): + return StructureHandler.from_data(self._raw_current_density.structure) + + def _read_current_densities(self): + return dict( + self._read_current_density(key) for key in self._raw_current_density + ) + + def _read_current_density(self, key=None): + key = key or self._raw_current_density.valid_indices[-1] + if key not in self._raw_current_density.valid_indices: + raise exception.IncorrectUsage( + f"Selection {key!r} is not available. " + f"Please use one of {list(self._raw_current_density.valid_indices)}." + ) + return f"current_{key}", self._raw_current_density[key].current_density[:].T + + +@quantity("current_density") +class CurrentDensity: + """Represents current density on the grid in the unit cell. + + A current density j is a vectorial quantity (j_x, j_y, j_z) on every grid point. + It describes how the current flows at every point in space. + + Examples + -------- + + First, we create some example data that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + To produce current density plots, please check the `to_contour` and `to_quiver` + functions for a more detailed documentation. + + To produce a contour plot: + + >>> calculation.current_density.to_contour("nmr", a=0) + Graph(series=[Contour(data=array([[...]]), ..., cut='a', ...)], ...) + + To produce a quiver plot: + + >>> calculation.current_density.to_quiver("nmr", c=0) + Graph(series=[Contour(data=array([[[...]]]), ..., cut='c', ...)], ...) + + For your own postprocessing, you can read the current density data into a Python dict: + + >>> calculation.current_density.read("nmr") + {'structure': {...}, 'current_x': array([[[[...]]]], ...), 'current_y': array([[[[...]]]], ...), 'current_z': array([[[[...]]]], ...)} + + You can inspect possible choices with: + + >>> calculation.current_density.selections("nmr") + {'current_density': ['nmr']} + + Please check the documentation of these methods for more details on how to use them + and which options they provide.""" + + def __init__(self, source, quantity_name="current_density"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_current_density): + return cls(source=DataSource(raw_current_density)) + + def _handler_factory(self, raw): + return CurrentDensityHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + CurrentDensityHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Read the current density and structural information into a Python dictionary. + + Returns + ------- + dict + Contains all available current density data as well as structural information. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + CurrentDensityHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) + def to_contour( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + """Generate a contour plot of current density. + + {plane} + + Parameters + ---------- + {parameters} + + Returns + ------- + Graph + A current density plot in the plane spanned by the 2 remaining lattice + vectors. + + Examples + -------- + + Cut a plane at the origin of the third lattice vector. + + >>> calculation.current_density.to_contour("nmr", c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calculation.current_density.to_contour("nmr", b=0.5, supercell=2) + + Take a slice along the first lattice vector and rotate it such that the normal + of the plane aligns with the x axis. + + >>> calculation.current_density.to_contour("nmr", a=0.3, normal="x") + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + CurrentDensityHandler.to_contour, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) + def to_quiver( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + """Generate a quiver plot of current density. + + {plane} + + Parameters + ---------- + {parameters} + + Returns + ------- + Graph + A quiver plot in the plane spanned by the 2 remaining lattice vectors. + + Examples + -------- + + Cut a plane at the origin of the third lattice vector. + + >>> calculation.current_density.to_quiver("nmr", c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calculation.current_density.to_quiver("nmr", b=0.5, supercell=2) + + Take a slice along the first lattice vector and rotate it such that the normal + of the plane aligns with the x axis. + + >>> calculation.current_density.to_quiver("nmr", a=0.3, normal="x") + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + CurrentDensityHandler.to_quiver, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + CurrentDensityHandler.from_data, + CurrentDensityHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/density.py b/src/py4vasp/_calculation/density.py index 181b47817..db6dd442c 100644 --- a/src/py4vasp/_calculation/density.py +++ b/src/py4vasp/_calculation/density.py @@ -1,728 +1,728 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy -import itertools -from typing import Optional, Union - -import numpy as np - -from py4vasp import _config, exception -from py4vasp import raw as raw_module -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation.dispatch import ( - DataSource, - merge_default, - merge_strings, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw import data as raw -from py4vasp._third_party import graph, view -from py4vasp._util import documentation, import_, index, select, slicing -from py4vasp._util.density import SliceArguments, Visualizer - -pretty = import_.optional("IPython.lib.pretty") - -_DEFAULT = 0 -_INTERNAL = "_density" -_COMPONENTS = { - 0: ["0", "unity", "sigma_0", "scalar", _INTERNAL], - 1: ["1", "sigma_x", "x", "sigma_1"], - 2: ["2", "sigma_y", "y", "sigma_2"], - 3: ["3", "sigma_z", "z", "sigma_3"], -} -_MAGNETIZATION = ("magnetization", "mag", "m") -_COMMON_PARAMETERS = f"""\ -{slicing.PARAMETERS} -supercell : int | np.ndarray | None = None - Replicate the contour plot periodically a given number of times. If you - provide two different numbers, the resulting cell will be the two remaining - lattice vectors multiplied by the specific number. -""" - - -def _join_with_emphasis(data): - emph_data = [f"*{x}*" for x in filter(lambda key: key != _INTERNAL, data)] - if len(data) < 3: - return " and ".join(emph_data) - emph_data.insert(-1, "and") - return ", ".join(emph_data) - - -class DensityHandler: - """Handler for density data — performs all data access and transformation.""" - - def __init__(self, raw_density: raw.Density, selection_name=None): - self._raw_density = raw_density - self._selection_name = selection_name - - @classmethod - def from_data( - cls, raw_density: raw.Density, selection_name=None - ) -> "DensityHandler": - return cls(raw_density, selection_name=selection_name) - - def __str__(self) -> str: - _raise_error_if_no_data(self._raw_density.charge) - grid = self._raw_density.charge.shape[1:] - raw_stoichiometry = self._raw_density.structure.stoichiometry - stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry) - if self._selection == "kinetic_energy": - name = "Kinetic energy" - elif self.is_nonpolarized(): - name = "Nonpolarized" - elif self.is_collinear(): - name = "Collinear" - else: - name = "Noncollinear" - return f"""{name} density: - structure: {pretty.pretty(stoichiometry)} - grid: {grid[2]}, {grid[1]}, {grid[0]}""" - - def to_dict(self) -> dict: - _raise_error_if_no_data(self._raw_density.charge) - result = {"structure": self._structure().to_dict()} - result.update(self._read_density()) - return result - - def selections(self) -> dict: - if self._raw_density.charge.is_none(): - return {} - if self.is_nonpolarized(): - components = [_COMPONENTS[0][_DEFAULT]] - elif self.is_collinear(): - components = [_COMPONENTS[0][_DEFAULT], _COMPONENTS[3][_DEFAULT]] - else: - components = [_COMPONENTS[i][_DEFAULT] for i in range(4)] - return {"component": components} - - def to_numpy(self): - return np.moveaxis(self._raw_density.charge, 0, -1).T - - def to_view( - self, - component: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - **user_options, - ) -> view.View: - _raise_error_if_no_data(self._raw_density.charge) - map_ = self._create_map() - selector = index.Selector({0: map_}, self._raw_density.charge) - component = component or _INTERNAL - tree = select.Tree.from_selection(component) - selections = list(self._filter_noncollinear_magnetization_from_selections(tree)) - structure_handler = self._structure() - viewer = structure_handler.to_view(supercell) - viewer.grid_scalars = [ - view.GridQuantity( - quantity=selector[sel].T[np.newaxis], - label=self._label(selector.label(sel)), - isosurfaces=self._grid_quantity_properties( - selector, sel, map_, user_options - ), - ) - for sel in selections - ] - return viewer - - def to_contour( - self, - component: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - map_ = self._create_map() - selector = index.Selector({0: map_}, self._raw_density.charge) - component = component or _INTERNAL - tree = select.Tree.from_selection(component) - selections = list(self._filter_noncollinear_magnetization_from_selections(tree)) - visualizer = Visualizer(self._structure()) - dataDict = { - (self._label(selector.label(sel)) or "charge"): selector[sel].T - for sel in selections - } - return visualizer.to_contour( - dataDict, SliceArguments(a, b, c, normal, supercell) - ) - - def to_quiver( - self, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - if self.is_collinear(): - data = self._raw_density.charge[1].T - else: - data = self.to_numpy()[1:] - visualizer = Visualizer(self._structure()) - dataDict = {(self._selection or "magnetization"): data} - return visualizer.to_quiver( - dataDict, SliceArguments(a, b, c, normal, supercell) - ) - - def is_nonpolarized(self): - return len(self._raw_density.charge) == 1 - - def is_collinear(self): - return len(self._raw_density.charge) == 2 - - def is_noncollinear(self): - return len(self._raw_density.charge) == 4 - - @property - def _selection(self): - selection_map = { - "kinetic_energy": "kinetic_energy", - "kinetic_energy_density": "kinetic_energy", - } - return selection_map.get(self._selection_name) - - def _structure(self): - return StructureHandler.from_data(self._raw_density.structure) - - def _read_density(self): - density = self.to_numpy() - if self._selection: - yield self._selection, density - else: - yield "charge", density[0] - if self.is_collinear(): - yield "magnetization", density[1] - elif self.is_noncollinear(): - yield "magnetization", density[1:] - - def _filter_noncollinear_magnetization_from_selections(self, tree): - if self._selection or not self.is_noncollinear(): - yield from tree.selections() - else: - filtered_selections = tree.selections(filter=set(_MAGNETIZATION)) - for filtered, unfiltered in zip(filtered_selections, tree.selections()): - if filtered != unfiltered and len(filtered) != 1: - _raise_component_not_specified_error(unfiltered) - yield filtered - - def _create_map(self): - map_ = { - choice: self._index_component(component) - for component, choices in _COMPONENTS.items() - for choice in choices - } - self._add_magnetization_for_charge_and_collinear(map_) - return map_ - - def _index_component(self, component): - if self.is_collinear(): - component = (0, 2, 3, 1)[component] - return component - - def _add_magnetization_for_charge_and_collinear(self, map_): - if self._selection or not self.is_collinear(): - return - for key in _MAGNETIZATION: - map_[key] = 1 - - def _grid_quantity_properties(self, selector, selection, map_, user_options): - component_label = selector.label(selection) - component = map_.get(component_label, -1) - return self._isosurfaces(component, **user_options) - - def _label(self, component_label): - if component_label == _INTERNAL: - return self._selection or "charge" - elif self._selection: - return f"{self._selection}({component_label})" - else: - return component_label - - def _isosurfaces(self, component, isolevel=0.2, color=None, opacity=0.6): - if self._use_symmetric_isosurface(component): - _raise_error_if_color_is_specified(color) - return [ - view.Isosurface(isolevel, _config.VASP_COLORS["blue"], opacity), - view.Isosurface(-isolevel, _config.VASP_COLORS["red"], opacity), - ] - else: - color = color or _config.VASP_COLORS["cyan"] - return [view.Isosurface(isolevel, color, opacity)] - - def _use_symmetric_isosurface(self, component): - if component > 0 and self.is_nonpolarized(): - _raise_is_nonpolarized_error() - if component > 1 and self.is_collinear(): - _raise_is_collinear_error() - return component > 0 - - -@quantity("density") -class Density(view.Mixin): - """This class accesses various densities (charge, magnetization, ...) of VASP. - - The charge density is one key quantity optimized by VASP. With this class you - can extract the final density and visualize it within the structure of the - system. For collinear calculations, one can also consider the magnetization - density. For noncollinear calculations, the magnetization density has three - components. One may also be interested in the kinetic energy density for - metaGGA calculations. - - Examples - -------- - - First, we create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - To produce density plots, please check the `to_contour` and `to_quiver` functions for - a more detailed documentation. - - To produce a contour plot: - - >>> calculation.density.to_contour(a=0) - Graph(series=[Contour(data=array([[...]]), ..., cut='a', ...)], ...) - - You can also visualize a 3d isosurface of the density: - - >>> calculation.density.plot() - View(elements=array([[...]]...), lattice_vectors=array([[[...]]]...), positions=array([[[...]]]...), grid_scalars=[GridQuantity(quantity=array([[[[...]]]]...), label='charge', isosurfaces=[Isosurface(...)])], ...) - - For your own postprocessing, you can read the band data into a Python dictionary: - - >>> calculation.density.read() - {'structure': ..., 'charge': array([[[...]]], ...)} - - Alternatively, obtain the density as a numpy array directly: - - >>> calculation.density.to_numpy() - array([[[[...]]]], ...) - - It is also possible to test for non-polarized, collinear, and noncollinear calculations - with: - - >>> calculation.density.is_nonpolarized() - True - >>> calculation.density.is_collinear() - False - >>> calculation.density.is_noncollinear() - False - - You can inspect possible selections with: - - >>> calculation.density.selections() - {'density': [...], 'component': ['0']} - - To produce a quiver plot for a noncollinear calculation: - - >>> from py4vasp import demo - >>> calculation_nc = demo.calculation(path, selection="noncollinear") - >>> calculation_nc.density.is_noncollinear() - True - >>> calculation_nc.density.to_quiver(c=0, supercell=2) - Graph(series=[Contour(data=array([[[...]]]), ..., cut='c', ...)], ...) - - Please check the documentation of these methods for more details on how to use them - and which options they provide. - """ - - def __init__(self, source, quantity_name="density", selection_name=None): - self._source = source - self._quantity_name = quantity_name - self._selection_name = selection_name - - @classmethod - def from_data(cls, raw_density): - return cls(source=DataSource(raw_density)) - - def __getitem__(self, selection_name) -> "Density": - new = copy.copy(self) - new._selection_name = selection_name - return new - - def _handler_factory(self, raw): - return DensityHandler.from_data(raw, selection_name=self._selection_name) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection or self._selection_name, - self._handler_factory, - DensityHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Read the density into a dictionary. - - Parameters - ---------- - selection : str - VASP computes different densities depending on the INCAR settings. With this - parameter, you can control which one of them is returned. Please use the - `selections` routine to get a list of all possible choices. - - Returns - ------- - dict - Contains the structure information as well as the density represented - on a grid in the unit cell. - """ - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - @documentation.format( - component0=_join_with_emphasis(_COMPONENTS[0]), - component1=_join_with_emphasis(_COMPONENTS[1]), - component2=_join_with_emphasis(_COMPONENTS[2]), - component3=_join_with_emphasis(_COMPONENTS[3]), - ) - def selections(self) -> dict: - """Returns possible densities VASP can produce along with all available components. - - In the dictionary, the key *density* lists all different densities you can access - from the VASP output provided you set the relevant INCAR tags. You can combine - any of these with any possible choice from the key *component* to further - specify the particular output you will receive. If you do not specify a *density* - or a *component* the other routines will default to the electronic charge and - the 0-th component. - - To nest density and component, please use parentheses, e.g. ``charge(1, 2)`` or - ``3(kinetic_energy)``. - - For convenience, py4vasp accepts the following aliases - - electronic charge density - *charge*, *n*, *charge_density*, and *electronic_charge_density* - - kinetic energy density - *kinetic_energy*, *kinetic_energy*, and *kinetic_energy_density* - - 0th component - {component0} - - 1st component - {component1} - - 2nd component - {component2} - - 3rd component - {component3} - - Returns - ------- - dict - Possible densities and components to pass as selection in other functions - on density. - - Notes - ----- - In the special case of collinear calculations, *magnetization*, *mag*, and *m* - are another alias for the 3rd component of the charge density. - - Examples - -------- - >>> calculation = py4vasp.Calculation.from_path(".") - >>> calculation.density.to_dict("n") - >>> calculation.density.plot("magnetization") - - Using synonyms and nesting - - >>> calculation.density.plot("n m(1,2) mag(sigma_z)") - """ - result = merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.selections, - ) - result["density"] = list(raw_module.selections(self._quantity_name)) - return result - - def to_numpy(self): - """Convert the density to a numpy array. - - The number of components is 1 for nonpolarized calculations, 2 for collinear - calculations, and 4 for noncollinear calculations. Each component is 3 - dimensional according to the grid VASP uses for the FFTs. - - Returns - ------- - np.ndarray - All components of the selected density. - """ - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.to_numpy, - ) - - def to_view( - self, - selection: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - **user_options, - ) -> view.View: - """Plot the selected density as a 3d isosurface within the structure. - - Parameters - ---------- - selection : str | None = None - Can be either *charge* or *magnetization*, depending on which quantity - should be visualized. For a noncollinear calculation, the density has - 4 components which can be represented in a 2x2 matrix. Specify the - component of the density in terms of the Pauli matrices: sigma_1, - sigma_2, sigma_3. - - supercell : int | np.ndarray | None = None - If present the data is replicated the specified number of times along each - direction. - - user_options : dict - Further arguments with keyword that get directly passed on to the - visualizer. Most importantly, you can set isolevel to adjust the - value at which the isosurface is drawn. - - Returns - ------- - View - Visualize an isosurface of the density within the 3d structure. - - Examples - -------- - - >>> calculation = py4vasp.Calculation.from_path(".") - - Plot an isosurface of the electronic charge density - - >>> calculation.density.plot() - - Plot isosurfaces for positive (blue) and negative (red) magnetization - of a spin-polarized calculation (ISPIN=2) - - >>> calculation.density.plot("m") - - Plot the isosurface for the third component of a noncollinear magnetization - - >>> calculation.density.plot("m(3)") - """ - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.to_view, - selection, - supercell=supercell, - **user_options, - ) - - @documentation.format(plane=slicing.PLANE, common_parameters=_COMMON_PARAMETERS) - def to_contour( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - """Generate a contour plot of the selected component of the density. - - {plane} - - Parameters - ---------- - selection : str | None = None - Select which component of the density you want to visualize. Please use the - `selections` method to get all available choices. - - {common_parameters} - - Returns - ------- - Graph - A contour plot in the plane spanned by the 2 remaining lattice vectors. - - - Examples - -------- - - Cut a plane through the magnetization density at the origin of the third lattice - vector. - - >>> calculation.density.to_contour("3", c=0) - - Replicate a plane in the middle of the second lattice vector 2 times in each - direction. - - >>> calculation.density.to_contour(b=0.5, supercell=2) - - Take a slice of the kinetic energy density along the first lattice vector and - rotate it such that the normal of the plane aligns with the x axis. - - >>> calculation.density.to_contour("kinetic_energy", a=0.3, normal="x") - """ - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.to_contour, - selection, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - @documentation.format(plane=slicing.PLANE, common_parameters=_COMMON_PARAMETERS) - def to_quiver( - self, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ) -> graph.Graph: - """Generate a quiver plot of magnetization density. - - {plane} - - For a collinear calculation, the magnetization density will be aligned with the - y axis of the plane. For noncollinear calculations, the magnetization density - is projected into the plane. - - Parameters - ---------- - {common_parameters} - - Returns - ------- - Graph - A quiver plot in the plane spanned by the 2 remaining lattice vectors. - - - Examples - -------- - - Cut a plane at the origin of the third lattice vector. - - >>> calculation.density.to_quiver(c=0) - - Replicate a plane in the middle of the second lattice vector 2 times in each - direction. - - >>> calculation.density.to_quiver(b=0.5, supercell=2) - - Take a slice of the spin components of the kinetic energy density along the - first lattice vector and rotate it such that the normal of the plane aligns with - the x axis. - - >>> calculation.density.to_quiver("kinetic_energy", a=0.3, normal="x") - """ - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.to_quiver, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - def is_nonpolarized(self, selection=None): - "Returns whether the density is not spin polarized." - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.is_nonpolarized, - ) - - def is_collinear(self, selection=None): - "Returns whether the density has a collinear magnetization." - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.is_collinear, - ) - - def is_noncollinear(self, selection=None): - "Returns whether the density has a noncollinear magnetization." - return merge_default( - self._source, - self._quantity_name, - self._selection_name, - self._handler_factory, - DensityHandler.is_noncollinear, - ) - - -def _raise_error_if_color_is_specified(color): - if color is not None: - msg = "Specifying the color of a magnetic isosurface is not implemented." - raise exception.NotImplemented(msg) - - -def _raise_component_not_specified_error(selec_tuple): - msg = ( - "Invalid selection: selection='" - + ", ".join(selec_tuple) - + "'. For a noncollinear calculation, the density has 4 components which can be represented in a 2x2 matrix. Specify the component of the density in terms of the Pauli matrices: sigma_1, sigma_2, sigma_3. E.g.: m(sigma_1)." - ) - raise exception.IncorrectUsage(msg) - - -def _raise_is_nonpolarized_error(): - msg = "Density does not contain magnetization. Please rerun VASP with ISPIN = 2 or LNONCOLLINEAR = T to obtain it." - raise exception.NoData(msg) - - -def _raise_is_collinear_error(): - msg = "Density does not contain noncollinear magnetization. Please rerun VASP with LNONCOLLINEAR = T to obtain it." - raise exception.NoData(msg) - - -def _raise_error_if_no_data(data): - if data.is_none(): - raise exception.NoData( - "Density data was not found. Note that the density information is written " - "on the demand to a different file (vaspwave.h5). Please make sure that " - "this file exists and LCHARGH5 = T is set in the INCAR file. Another " - 'common issue is when you create `Calculation.from_file("vaspout.h5")` ' - "because this will overwrite the default file behavior." - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy +import itertools +from typing import Optional, Union + +import numpy as np + +from py4vasp import _config, exception +from py4vasp import raw as raw_module +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw import data as raw +from py4vasp._third_party import graph, view +from py4vasp._util import documentation, import_, index, select, slicing +from py4vasp._util.density import SliceArguments, Visualizer + +pretty = import_.optional("IPython.lib.pretty") + +_DEFAULT = 0 +_INTERNAL = "_density" +_COMPONENTS = { + 0: ["0", "unity", "sigma_0", "scalar", _INTERNAL], + 1: ["1", "sigma_x", "x", "sigma_1"], + 2: ["2", "sigma_y", "y", "sigma_2"], + 3: ["3", "sigma_z", "z", "sigma_3"], +} +_MAGNETIZATION = ("magnetization", "mag", "m") +_COMMON_PARAMETERS = f"""\ +{slicing.PARAMETERS} +supercell : int | np.ndarray | None = None + Replicate the contour plot periodically a given number of times. If you + provide two different numbers, the resulting cell will be the two remaining + lattice vectors multiplied by the specific number. +""" + + +def _join_with_emphasis(data): + emph_data = [f"*{x}*" for x in filter(lambda key: key != _INTERNAL, data)] + if len(data) < 3: + return " and ".join(emph_data) + emph_data.insert(-1, "and") + return ", ".join(emph_data) + + +class DensityHandler: + """Handler for density data — performs all data access and transformation.""" + + def __init__(self, raw_density: raw.Density, selection_name=None): + self._raw_density = raw_density + self._selection_name = selection_name + + @classmethod + def from_data( + cls, raw_density: raw.Density, selection_name=None + ) -> "DensityHandler": + return cls(raw_density, selection_name=selection_name) + + def __str__(self) -> str: + _raise_error_if_no_data(self._raw_density.charge) + grid = self._raw_density.charge.shape[1:] + raw_stoichiometry = self._raw_density.structure.stoichiometry + stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry) + if self._selection == "kinetic_energy": + name = "Kinetic energy" + elif self.is_nonpolarized(): + name = "Nonpolarized" + elif self.is_collinear(): + name = "Collinear" + else: + name = "Noncollinear" + return f"""{name} density: + structure: {pretty.pretty(stoichiometry)} + grid: {grid[2]}, {grid[1]}, {grid[0]}""" + + def to_dict(self) -> dict: + _raise_error_if_no_data(self._raw_density.charge) + result = {"structure": self._structure().to_dict()} + result.update(self._read_density()) + return result + + def selections(self) -> dict: + if self._raw_density.charge.is_none(): + return {} + if self.is_nonpolarized(): + components = [_COMPONENTS[0][_DEFAULT]] + elif self.is_collinear(): + components = [_COMPONENTS[0][_DEFAULT], _COMPONENTS[3][_DEFAULT]] + else: + components = [_COMPONENTS[i][_DEFAULT] for i in range(4)] + return {"component": components} + + def to_numpy(self): + return np.moveaxis(self._raw_density.charge, 0, -1).T + + def to_view( + self, + component: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + **user_options, + ) -> view.View: + _raise_error_if_no_data(self._raw_density.charge) + map_ = self._create_map() + selector = index.Selector({0: map_}, self._raw_density.charge) + component = component or _INTERNAL + tree = select.Tree.from_selection(component) + selections = list(self._filter_noncollinear_magnetization_from_selections(tree)) + structure_handler = self._structure() + viewer = structure_handler.to_view(supercell) + viewer.grid_scalars = [ + view.GridQuantity( + quantity=selector[sel].T[np.newaxis], + label=self._label(selector.label(sel)), + isosurfaces=self._grid_quantity_properties( + selector, sel, map_, user_options + ), + ) + for sel in selections + ] + return viewer + + def to_contour( + self, + component: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + map_ = self._create_map() + selector = index.Selector({0: map_}, self._raw_density.charge) + component = component or _INTERNAL + tree = select.Tree.from_selection(component) + selections = list(self._filter_noncollinear_magnetization_from_selections(tree)) + visualizer = Visualizer(self._structure()) + dataDict = { + (self._label(selector.label(sel)) or "charge"): selector[sel].T + for sel in selections + } + return visualizer.to_contour( + dataDict, SliceArguments(a, b, c, normal, supercell) + ) + + def to_quiver( + self, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + if self.is_collinear(): + data = self._raw_density.charge[1].T + else: + data = self.to_numpy()[1:] + visualizer = Visualizer(self._structure()) + dataDict = {(self._selection or "magnetization"): data} + return visualizer.to_quiver( + dataDict, SliceArguments(a, b, c, normal, supercell) + ) + + def is_nonpolarized(self): + return len(self._raw_density.charge) == 1 + + def is_collinear(self): + return len(self._raw_density.charge) == 2 + + def is_noncollinear(self): + return len(self._raw_density.charge) == 4 + + @property + def _selection(self): + selection_map = { + "kinetic_energy": "kinetic_energy", + "kinetic_energy_density": "kinetic_energy", + } + return selection_map.get(self._selection_name) + + def _structure(self): + return StructureHandler.from_data(self._raw_density.structure) + + def _read_density(self): + density = self.to_numpy() + if self._selection: + yield self._selection, density + else: + yield "charge", density[0] + if self.is_collinear(): + yield "magnetization", density[1] + elif self.is_noncollinear(): + yield "magnetization", density[1:] + + def _filter_noncollinear_magnetization_from_selections(self, tree): + if self._selection or not self.is_noncollinear(): + yield from tree.selections() + else: + filtered_selections = tree.selections(filter=set(_MAGNETIZATION)) + for filtered, unfiltered in zip(filtered_selections, tree.selections()): + if filtered != unfiltered and len(filtered) != 1: + _raise_component_not_specified_error(unfiltered) + yield filtered + + def _create_map(self): + map_ = { + choice: self._index_component(component) + for component, choices in _COMPONENTS.items() + for choice in choices + } + self._add_magnetization_for_charge_and_collinear(map_) + return map_ + + def _index_component(self, component): + if self.is_collinear(): + component = (0, 2, 3, 1)[component] + return component + + def _add_magnetization_for_charge_and_collinear(self, map_): + if self._selection or not self.is_collinear(): + return + for key in _MAGNETIZATION: + map_[key] = 1 + + def _grid_quantity_properties(self, selector, selection, map_, user_options): + component_label = selector.label(selection) + component = map_.get(component_label, -1) + return self._isosurfaces(component, **user_options) + + def _label(self, component_label): + if component_label == _INTERNAL: + return self._selection or "charge" + elif self._selection: + return f"{self._selection}({component_label})" + else: + return component_label + + def _isosurfaces(self, component, isolevel=0.2, color=None, opacity=0.6): + if self._use_symmetric_isosurface(component): + _raise_error_if_color_is_specified(color) + return [ + view.Isosurface(isolevel, _config.VASP_COLORS["blue"], opacity), + view.Isosurface(-isolevel, _config.VASP_COLORS["red"], opacity), + ] + else: + color = color or _config.VASP_COLORS["cyan"] + return [view.Isosurface(isolevel, color, opacity)] + + def _use_symmetric_isosurface(self, component): + if component > 0 and self.is_nonpolarized(): + _raise_is_nonpolarized_error() + if component > 1 and self.is_collinear(): + _raise_is_collinear_error() + return component > 0 + + +@quantity("density") +class Density(view.Mixin): + """This class accesses various densities (charge, magnetization, ...) of VASP. + + The charge density is one key quantity optimized by VASP. With this class you + can extract the final density and visualize it within the structure of the + system. For collinear calculations, one can also consider the magnetization + density. For noncollinear calculations, the magnetization density has three + components. One may also be interested in the kinetic energy density for + metaGGA calculations. + + Examples + -------- + + First, we create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + To produce density plots, please check the `to_contour` and `to_quiver` functions for + a more detailed documentation. + + To produce a contour plot: + + >>> calculation.density.to_contour(a=0) + Graph(series=[Contour(data=array([[...]]), ..., cut='a', ...)], ...) + + You can also visualize a 3d isosurface of the density: + + >>> calculation.density.plot() + View(elements=array([[...]]...), lattice_vectors=array([[[...]]]...), positions=array([[[...]]]...), grid_scalars=[GridQuantity(quantity=array([[[[...]]]]...), label='charge', isosurfaces=[Isosurface(...)])], ...) + + For your own postprocessing, you can read the band data into a Python dictionary: + + >>> calculation.density.read() + {'structure': ..., 'charge': array([[[...]]], ...)} + + Alternatively, obtain the density as a numpy array directly: + + >>> calculation.density.to_numpy() + array([[[[...]]]], ...) + + It is also possible to test for non-polarized, collinear, and noncollinear calculations + with: + + >>> calculation.density.is_nonpolarized() + True + >>> calculation.density.is_collinear() + False + >>> calculation.density.is_noncollinear() + False + + You can inspect possible selections with: + + >>> calculation.density.selections() + {'density': [...], 'component': ['0']} + + To produce a quiver plot for a noncollinear calculation: + + >>> from py4vasp import demo + >>> calculation_nc = demo.calculation(path, selection="noncollinear") + >>> calculation_nc.density.is_noncollinear() + True + >>> calculation_nc.density.to_quiver(c=0, supercell=2) + Graph(series=[Contour(data=array([[[...]]]), ..., cut='c', ...)], ...) + + Please check the documentation of these methods for more details on how to use them + and which options they provide. + """ + + def __init__(self, source, quantity_name="density", selection_name=None): + self._source = source + self._quantity_name = quantity_name + self._selection_name = selection_name + + @classmethod + def from_data(cls, raw_density): + return cls(source=DataSource(raw_density)) + + def __getitem__(self, selection_name) -> "Density": + new = copy.copy(self) + new._selection_name = selection_name + return new + + def _handler_factory(self, raw): + return DensityHandler.from_data(raw, selection_name=self._selection_name) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection or self._selection_name, + self._handler_factory, + DensityHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Read the density into a dictionary. + + Parameters + ---------- + selection : str + VASP computes different densities depending on the INCAR settings. With this + parameter, you can control which one of them is returned. Please use the + `selections` routine to get a list of all possible choices. + + Returns + ------- + dict + Contains the structure information as well as the density represented + on a grid in the unit cell. + """ + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + @documentation.format( + component0=_join_with_emphasis(_COMPONENTS[0]), + component1=_join_with_emphasis(_COMPONENTS[1]), + component2=_join_with_emphasis(_COMPONENTS[2]), + component3=_join_with_emphasis(_COMPONENTS[3]), + ) + def selections(self) -> dict: + """Returns possible densities VASP can produce along with all available components. + + In the dictionary, the key *density* lists all different densities you can access + from the VASP output provided you set the relevant INCAR tags. You can combine + any of these with any possible choice from the key *component* to further + specify the particular output you will receive. If you do not specify a *density* + or a *component* the other routines will default to the electronic charge and + the 0-th component. + + To nest density and component, please use parentheses, e.g. ``charge(1, 2)`` or + ``3(kinetic_energy)``. + + For convenience, py4vasp accepts the following aliases + + electronic charge density + *charge*, *n*, *charge_density*, and *electronic_charge_density* + + kinetic energy density + *kinetic_energy*, *kinetic_energy*, and *kinetic_energy_density* + + 0th component + {component0} + + 1st component + {component1} + + 2nd component + {component2} + + 3rd component + {component3} + + Returns + ------- + dict + Possible densities and components to pass as selection in other functions + on density. + + Notes + ----- + In the special case of collinear calculations, *magnetization*, *mag*, and *m* + are another alias for the 3rd component of the charge density. + + Examples + -------- + >>> calculation = py4vasp.Calculation.from_path(".") + >>> calculation.density.to_dict("n") + >>> calculation.density.plot("magnetization") + + Using synonyms and nesting + + >>> calculation.density.plot("n m(1,2) mag(sigma_z)") + """ + result = merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.selections, + ) + result["density"] = list(raw_module.selections(self._quantity_name)) + return result + + def to_numpy(self): + """Convert the density to a numpy array. + + The number of components is 1 for nonpolarized calculations, 2 for collinear + calculations, and 4 for noncollinear calculations. Each component is 3 + dimensional according to the grid VASP uses for the FFTs. + + Returns + ------- + np.ndarray + All components of the selected density. + """ + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.to_numpy, + ) + + def to_view( + self, + selection: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + **user_options, + ) -> view.View: + """Plot the selected density as a 3d isosurface within the structure. + + Parameters + ---------- + selection : str | None = None + Can be either *charge* or *magnetization*, depending on which quantity + should be visualized. For a noncollinear calculation, the density has + 4 components which can be represented in a 2x2 matrix. Specify the + component of the density in terms of the Pauli matrices: sigma_1, + sigma_2, sigma_3. + + supercell : int | np.ndarray | None = None + If present the data is replicated the specified number of times along each + direction. + + user_options : dict + Further arguments with keyword that get directly passed on to the + visualizer. Most importantly, you can set isolevel to adjust the + value at which the isosurface is drawn. + + Returns + ------- + View + Visualize an isosurface of the density within the 3d structure. + + Examples + -------- + + >>> calculation = py4vasp.Calculation.from_path(".") + + Plot an isosurface of the electronic charge density + + >>> calculation.density.plot() + + Plot isosurfaces for positive (blue) and negative (red) magnetization + of a spin-polarized calculation (ISPIN=2) + + >>> calculation.density.plot("m") + + Plot the isosurface for the third component of a noncollinear magnetization + + >>> calculation.density.plot("m(3)") + """ + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.to_view, + selection, + supercell=supercell, + **user_options, + ) + + @documentation.format(plane=slicing.PLANE, common_parameters=_COMMON_PARAMETERS) + def to_contour( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + """Generate a contour plot of the selected component of the density. + + {plane} + + Parameters + ---------- + selection : str | None = None + Select which component of the density you want to visualize. Please use the + `selections` method to get all available choices. + + {common_parameters} + + Returns + ------- + Graph + A contour plot in the plane spanned by the 2 remaining lattice vectors. + + + Examples + -------- + + Cut a plane through the magnetization density at the origin of the third lattice + vector. + + >>> calculation.density.to_contour("3", c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calculation.density.to_contour(b=0.5, supercell=2) + + Take a slice of the kinetic energy density along the first lattice vector and + rotate it such that the normal of the plane aligns with the x axis. + + >>> calculation.density.to_contour("kinetic_energy", a=0.3, normal="x") + """ + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.to_contour, + selection, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + @documentation.format(plane=slicing.PLANE, common_parameters=_COMMON_PARAMETERS) + def to_quiver( + self, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ) -> graph.Graph: + """Generate a quiver plot of magnetization density. + + {plane} + + For a collinear calculation, the magnetization density will be aligned with the + y axis of the plane. For noncollinear calculations, the magnetization density + is projected into the plane. + + Parameters + ---------- + {common_parameters} + + Returns + ------- + Graph + A quiver plot in the plane spanned by the 2 remaining lattice vectors. + + + Examples + -------- + + Cut a plane at the origin of the third lattice vector. + + >>> calculation.density.to_quiver(c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calculation.density.to_quiver(b=0.5, supercell=2) + + Take a slice of the spin components of the kinetic energy density along the + first lattice vector and rotate it such that the normal of the plane aligns with + the x axis. + + >>> calculation.density.to_quiver("kinetic_energy", a=0.3, normal="x") + """ + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.to_quiver, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + def is_nonpolarized(self, selection=None): + "Returns whether the density is not spin polarized." + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.is_nonpolarized, + ) + + def is_collinear(self, selection=None): + "Returns whether the density has a collinear magnetization." + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.is_collinear, + ) + + def is_noncollinear(self, selection=None): + "Returns whether the density has a noncollinear magnetization." + return merge_default( + self._source, + self._quantity_name, + self._selection_name, + self._handler_factory, + DensityHandler.is_noncollinear, + ) + + +def _raise_error_if_color_is_specified(color): + if color is not None: + msg = "Specifying the color of a magnetic isosurface is not implemented." + raise exception.NotImplemented(msg) + + +def _raise_component_not_specified_error(selec_tuple): + msg = ( + "Invalid selection: selection='" + + ", ".join(selec_tuple) + + "'. For a noncollinear calculation, the density has 4 components which can be represented in a 2x2 matrix. Specify the component of the density in terms of the Pauli matrices: sigma_1, sigma_2, sigma_3. E.g.: m(sigma_1)." + ) + raise exception.IncorrectUsage(msg) + + +def _raise_is_nonpolarized_error(): + msg = "Density does not contain magnetization. Please rerun VASP with ISPIN = 2 or LNONCOLLINEAR = T to obtain it." + raise exception.NoData(msg) + + +def _raise_is_collinear_error(): + msg = "Density does not contain noncollinear magnetization. Please rerun VASP with LNONCOLLINEAR = T to obtain it." + raise exception.NoData(msg) + + +def _raise_error_if_no_data(data): + if data.is_none(): + raise exception.NoData( + "Density data was not found. Note that the density information is written " + "on the demand to a different file (vaspwave.h5). Please make sure that " + "this file exists and LCHARGH5 = T is set in the INCAR file. Another " + 'common issue is when you create `Calculation.from_file("vaspout.h5")` ' + "because this will overwrite the default file behavior." + ) diff --git a/src/py4vasp/_calculation/dielectric_function.py b/src/py4vasp/_calculation/dielectric_function.py index 5ddc1f0b1..f8d750f70 100644 --- a/src/py4vasp/_calculation/dielectric_function.py +++ b/src/py4vasp/_calculation/dielectric_function.py @@ -1,365 +1,365 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) - -import numpy as np - -from py4vasp import raw -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import DielectricFunction_DB -from py4vasp._third_party import graph -from py4vasp._util import check, convert, index, select - - -class DielectricFunctionHandler: - """Handler for the dielectric_function quantity. Works with exactly one raw.DielectricFunction object.""" - - def __init__(self, raw_dielectric_function: raw.DielectricFunction): - self._raw_dielectric_function = raw_dielectric_function - - @classmethod - def from_data( - cls, raw_dielectric_function: raw.DielectricFunction - ) -> "DielectricFunctionHandler": - return cls(raw_dielectric_function) - - def to_dict(self) -> dict: - """Read the data into a dictionary. - - Returns - ------- - dict - Contains the energies at which the dielectric function was evaluated - and the dielectric tensor (3x3 matrix) at these energies. - """ - data = convert.to_complex( - np.array(self._raw_dielectric_function.dielectric_function) - ) - return { - "energies": self._raw_dielectric_function.energies[:], - "dielectric_function": data, - **self._add_current_current_if_available(), - **self._add_q_point_if_available(), - } - - def to_database(self) -> dict: - """Serialize dielectric function data for database storage.""" - return DielectricFunction_DB( - energy_min=( - float(np.min(self._raw_dielectric_function.energies[:])) - if not check.is_none(self._raw_dielectric_function.energies) - else None - ), - energy_max=( - float(np.max(self._raw_dielectric_function.energies[:])) - if not check.is_none(self._raw_dielectric_function.energies) - else None - ), - ) - - def to_graph(self, selection=None) -> graph.Graph: - """Read the data and generate a figure with the selected directions. - - Parameters - ---------- - selection : str - Specify along which directions and which components of the dielectric - function you want to plot. Defaults to *isotropic* and both the real - and the complex part. You can use the `selections` routine if you are - not sure which options are available. - - Returns - ------- - Graph - figure containing the dielectric function for the selected - directions and components. - """ - selection = self._replace_complex_labels(selection or "") - return graph.Graph( - series=self._make_series(selection), - xlabel="Energy (eV)", - ylabel="dielectric function ϵ", - ) - - def selections(self) -> dict: - """Returns a dictionary of possible selections for component, direction, and complex value.""" - complex_selections = {"complex": ["real", "Re", "imag", "Im"]} - if not self._has_tensor_data(): - return complex_selections - components = ( - ["density", "current"] if self._has_current_component() else ["density"] - ) - return { - "components": components, - "directions": [key for key in self._init_directions_dict() if key], - **complex_selections, - } - - def __str__(self) -> str: - energies = self._raw_dielectric_function.energies - header = f"""\ -dielectric function: - energies: [{energies[0]:0.2f}, {energies[-1]:0.2f}] {len(energies)} points""" - if self._has_tensor_data(): - footer = "directions: isotropic, xx, yy, zz, xy, yz, xz" - else: - qpoint_label = ", ".join( - f"{q:0.3f}" for q in self._raw_dielectric_function.q_point - ) - footer = f"q-point: [{qpoint_label}]" - if self._has_current_component(): - return f"""\ -{header} - components: density, current - {footer}""" - else: - return f"""\ -{header} - {footer}""" - - def _add_current_current_if_available(self): - if self._has_current_component(): - data = convert.to_complex( - np.array(self._raw_dielectric_function.current_current) - ) - return {"current_current": data} - else: - return {} - - def _has_current_component(self): - return not check.is_none(self._raw_dielectric_function.current_current) - - def _add_q_point_if_available(self): - if self._has_q_point(): - return {"q_point": self._raw_dielectric_function.q_point[:]} - else: - return {} - - def _has_q_point(self): - return not check.is_none(self._raw_dielectric_function.q_point) - - def _replace_complex_labels(self, selection): - selection = selection.replace("real", "Re") - return selection.replace("imaginary", "Im").replace("imag", "Im") - - def _make_series(self, selection): - energies = self._raw_dielectric_function.energies[:] - selector = self._make_selector() - return [ - graph.Series( - energies, selector[selection], self._create_label(selector, selection) - ) - for selection in self._generate_selections(selection) - ] - - def _make_selector(self): - if self._has_tensor_data(): - maps = { - 3: self._init_complex_dict(), - 0: self._init_components_dict(), - 1: self._init_directions_dict(), - } - else: - maps = { - 1: self._init_complex_dict(), - } - return index.Selector(maps, self._get_data(), reduction=np.average) - - def _init_components_dict(self): - return {None: 0, "density": 0, "current": 1} - - def _init_directions_dict(self): - return { - None: [0, 4, 8], - "isotropic": [0, 4, 8], - "xx": 0, - "yy": 4, - "zz": 8, - "xy": [1, 3], - "xz": [2, 6], - "yz": [5, 7], - } - - def _init_complex_dict(self): - return {"Re": 0, "Im": 1} - - def _get_data(self): - *_, number_points, complex_ = ( - self._raw_dielectric_function.dielectric_function.shape - ) - if self._has_current_component(): - new_shape = (9, number_points, complex_) - density = np.reshape( - self._raw_dielectric_function.dielectric_function, new_shape - ) - current = np.reshape( - self._raw_dielectric_function.current_current, new_shape - ) - return np.array([density, current]) - elif self._has_tensor_data(): - new_shape = (1, 9, number_points, complex_) - return np.reshape( - self._raw_dielectric_function.dielectric_function, new_shape - ) - else: - return self._raw_dielectric_function.dielectric_function - - def _create_label(self, selector, selection): - if self._has_tensor_data(): - return selector.label(selection) - else: - q_point_label = ",".join( - str(convert.Fraction(q)) for q in self._raw_dielectric_function.q_point - ) - return f"{selector.label(selection)}_q=[{q_point_label}]" - - def _has_tensor_data(self): - return self._raw_dielectric_function.dielectric_function.ndim == 4 - - def _generate_selections(self, selection): - tree = select.Tree.from_selection(selection) - for selection in tree.selections(): - if not self._component_selected(selection): - selection = selection + ("density",) - if self._complex_selected(selection): - yield selection - else: - yield selection + ("Re",) - yield selection + ("Im",) - - def _component_selected(self, selection): - if self._has_current_component(): - return select.contains(selection, "density") or select.contains( - selection, "current" - ) - else: - return True - - def _complex_selected(self, selection): - return select.contains(selection, "Re") or select.contains(selection, "Im") - - -@quantity("dielectric_function") -class DielectricFunction(graph.Mixin): - """The dielectric function describes the material response to an electric field. - - The dielectric function is a fundamental concept that describes how a material - responds to an external electric field. It is a frequency-dependent complex-valued - 3x3 matrix that relates the polarization of a material to the applied electric - field. The dielectric function is essential in understanding optical properties, - such as refractive index and absorption. - - There are many different ways to compute dielectric functions with VASP. This - class provides a common interface to all of them. You can pass a `selection` - argument to any of the methods of this class to select which dielectric function - you are interested in. Please make sure the INCAR file you use is compatible - with the setup. - - The 3x3 matrix is symmetric so for the plotting routines, py4vasp uses only the - six distinct components (xx, yy, zz, xy, xz, yz). The default is the isotropic - dielectric function but you can also select specific components by providing - one of the six components as selection. - """ - - def __init__(self, source, quantity_name: str = "dielectric_function"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data( - cls, raw_dielectric_function: raw.DielectricFunction - ) -> "DielectricFunction": - """Create a DielectricFunction dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_dielectric_function)) - - @property - def path(self): - """Returns the path from which the output is obtained.""" - return self._path - - def _handler_factory(self, raw_data): - return DielectricFunctionHandler.from_data(raw_data) - - def read(self, selection: str | None = None) -> dict: - """Read the data into a dictionary. - - Returns - ------- - dict - Contains the energies at which the dielectric function was evaluated - and the dielectric tensor (3x3 matrix) at these energies. - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DielectricFunctionHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Public alias for read(). Check that method for examples and optional arguments.""" - return self.read(selection=selection) - - def to_graph(self, selection: str | None = None) -> graph.Graph: - """Read the data and generate a figure with the selected directions. - - Parameters - ---------- - selection : str - Specify along which directions and which components of the dielectric - function you want to plot. Defaults to *isotropic* and both the real - and the complex part. You can use the `selections` routine if you are - not sure which options are available. - - Returns - ------- - Graph - figure containing the dielectric function for the selected - directions and components. - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DielectricFunctionHandler.to_graph, - ) - - def selections(self, selection: str | None = None) -> dict: - """Returns a dictionary of possible selections for component, direction, and complex value.""" - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DielectricFunctionHandler.selections, - ) - - def __str__(self, selection: str | None = None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DielectricFunctionHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - DielectricFunctionHandler.from_data, - DielectricFunctionHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import numpy as np + +from py4vasp import raw +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import DielectricFunction_DB +from py4vasp._third_party import graph +from py4vasp._util import check, convert, index, select + + +class DielectricFunctionHandler: + """Handler for the dielectric_function quantity. Works with exactly one raw.DielectricFunction object.""" + + def __init__(self, raw_dielectric_function: raw.DielectricFunction): + self._raw_dielectric_function = raw_dielectric_function + + @classmethod + def from_data( + cls, raw_dielectric_function: raw.DielectricFunction + ) -> "DielectricFunctionHandler": + return cls(raw_dielectric_function) + + def to_dict(self) -> dict: + """Read the data into a dictionary. + + Returns + ------- + dict + Contains the energies at which the dielectric function was evaluated + and the dielectric tensor (3x3 matrix) at these energies. + """ + data = convert.to_complex( + np.array(self._raw_dielectric_function.dielectric_function) + ) + return { + "energies": self._raw_dielectric_function.energies[:], + "dielectric_function": data, + **self._add_current_current_if_available(), + **self._add_q_point_if_available(), + } + + def to_database(self) -> dict: + """Serialize dielectric function data for database storage.""" + return DielectricFunction_DB( + energy_min=( + float(np.min(self._raw_dielectric_function.energies[:])) + if not check.is_none(self._raw_dielectric_function.energies) + else None + ), + energy_max=( + float(np.max(self._raw_dielectric_function.energies[:])) + if not check.is_none(self._raw_dielectric_function.energies) + else None + ), + ) + + def to_graph(self, selection=None) -> graph.Graph: + """Read the data and generate a figure with the selected directions. + + Parameters + ---------- + selection : str + Specify along which directions and which components of the dielectric + function you want to plot. Defaults to *isotropic* and both the real + and the complex part. You can use the `selections` routine if you are + not sure which options are available. + + Returns + ------- + Graph + figure containing the dielectric function for the selected + directions and components. + """ + selection = self._replace_complex_labels(selection or "") + return graph.Graph( + series=self._make_series(selection), + xlabel="Energy (eV)", + ylabel="dielectric function ϵ", + ) + + def selections(self) -> dict: + """Returns a dictionary of possible selections for component, direction, and complex value.""" + complex_selections = {"complex": ["real", "Re", "imag", "Im"]} + if not self._has_tensor_data(): + return complex_selections + components = ( + ["density", "current"] if self._has_current_component() else ["density"] + ) + return { + "components": components, + "directions": [key for key in self._init_directions_dict() if key], + **complex_selections, + } + + def __str__(self) -> str: + energies = self._raw_dielectric_function.energies + header = f"""\ +dielectric function: + energies: [{energies[0]:0.2f}, {energies[-1]:0.2f}] {len(energies)} points""" + if self._has_tensor_data(): + footer = "directions: isotropic, xx, yy, zz, xy, yz, xz" + else: + qpoint_label = ", ".join( + f"{q:0.3f}" for q in self._raw_dielectric_function.q_point + ) + footer = f"q-point: [{qpoint_label}]" + if self._has_current_component(): + return f"""\ +{header} + components: density, current + {footer}""" + else: + return f"""\ +{header} + {footer}""" + + def _add_current_current_if_available(self): + if self._has_current_component(): + data = convert.to_complex( + np.array(self._raw_dielectric_function.current_current) + ) + return {"current_current": data} + else: + return {} + + def _has_current_component(self): + return not check.is_none(self._raw_dielectric_function.current_current) + + def _add_q_point_if_available(self): + if self._has_q_point(): + return {"q_point": self._raw_dielectric_function.q_point[:]} + else: + return {} + + def _has_q_point(self): + return not check.is_none(self._raw_dielectric_function.q_point) + + def _replace_complex_labels(self, selection): + selection = selection.replace("real", "Re") + return selection.replace("imaginary", "Im").replace("imag", "Im") + + def _make_series(self, selection): + energies = self._raw_dielectric_function.energies[:] + selector = self._make_selector() + return [ + graph.Series( + energies, selector[selection], self._create_label(selector, selection) + ) + for selection in self._generate_selections(selection) + ] + + def _make_selector(self): + if self._has_tensor_data(): + maps = { + 3: self._init_complex_dict(), + 0: self._init_components_dict(), + 1: self._init_directions_dict(), + } + else: + maps = { + 1: self._init_complex_dict(), + } + return index.Selector(maps, self._get_data(), reduction=np.average) + + def _init_components_dict(self): + return {None: 0, "density": 0, "current": 1} + + def _init_directions_dict(self): + return { + None: [0, 4, 8], + "isotropic": [0, 4, 8], + "xx": 0, + "yy": 4, + "zz": 8, + "xy": [1, 3], + "xz": [2, 6], + "yz": [5, 7], + } + + def _init_complex_dict(self): + return {"Re": 0, "Im": 1} + + def _get_data(self): + *_, number_points, complex_ = ( + self._raw_dielectric_function.dielectric_function.shape + ) + if self._has_current_component(): + new_shape = (9, number_points, complex_) + density = np.reshape( + self._raw_dielectric_function.dielectric_function, new_shape + ) + current = np.reshape( + self._raw_dielectric_function.current_current, new_shape + ) + return np.array([density, current]) + elif self._has_tensor_data(): + new_shape = (1, 9, number_points, complex_) + return np.reshape( + self._raw_dielectric_function.dielectric_function, new_shape + ) + else: + return self._raw_dielectric_function.dielectric_function + + def _create_label(self, selector, selection): + if self._has_tensor_data(): + return selector.label(selection) + else: + q_point_label = ",".join( + str(convert.Fraction(q)) for q in self._raw_dielectric_function.q_point + ) + return f"{selector.label(selection)}_q=[{q_point_label}]" + + def _has_tensor_data(self): + return self._raw_dielectric_function.dielectric_function.ndim == 4 + + def _generate_selections(self, selection): + tree = select.Tree.from_selection(selection) + for selection in tree.selections(): + if not self._component_selected(selection): + selection = selection + ("density",) + if self._complex_selected(selection): + yield selection + else: + yield selection + ("Re",) + yield selection + ("Im",) + + def _component_selected(self, selection): + if self._has_current_component(): + return select.contains(selection, "density") or select.contains( + selection, "current" + ) + else: + return True + + def _complex_selected(self, selection): + return select.contains(selection, "Re") or select.contains(selection, "Im") + + +@quantity("dielectric_function") +class DielectricFunction(graph.Mixin): + """The dielectric function describes the material response to an electric field. + + The dielectric function is a fundamental concept that describes how a material + responds to an external electric field. It is a frequency-dependent complex-valued + 3x3 matrix that relates the polarization of a material to the applied electric + field. The dielectric function is essential in understanding optical properties, + such as refractive index and absorption. + + There are many different ways to compute dielectric functions with VASP. This + class provides a common interface to all of them. You can pass a `selection` + argument to any of the methods of this class to select which dielectric function + you are interested in. Please make sure the INCAR file you use is compatible + with the setup. + + The 3x3 matrix is symmetric so for the plotting routines, py4vasp uses only the + six distinct components (xx, yy, zz, xy, xz, yz). The default is the isotropic + dielectric function but you can also select specific components by providing + one of the six components as selection. + """ + + def __init__(self, source, quantity_name: str = "dielectric_function"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data( + cls, raw_dielectric_function: raw.DielectricFunction + ) -> "DielectricFunction": + """Create a DielectricFunction dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_dielectric_function)) + + @property + def path(self): + """Returns the path from which the output is obtained.""" + return self._path + + def _handler_factory(self, raw_data): + return DielectricFunctionHandler.from_data(raw_data) + + def read(self, selection: str | None = None) -> dict: + """Read the data into a dictionary. + + Returns + ------- + dict + Contains the energies at which the dielectric function was evaluated + and the dielectric tensor (3x3 matrix) at these energies. + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DielectricFunctionHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Public alias for read(). Check that method for examples and optional arguments.""" + return self.read(selection=selection) + + def to_graph(self, selection: str | None = None) -> graph.Graph: + """Read the data and generate a figure with the selected directions. + + Parameters + ---------- + selection : str + Specify along which directions and which components of the dielectric + function you want to plot. Defaults to *isotropic* and both the real + and the complex part. You can use the `selections` routine if you are + not sure which options are available. + + Returns + ------- + Graph + figure containing the dielectric function for the selected + directions and components. + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DielectricFunctionHandler.to_graph, + ) + + def selections(self, selection: str | None = None) -> dict: + """Returns a dictionary of possible selections for component, direction, and complex value.""" + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DielectricFunctionHandler.selections, + ) + + def __str__(self, selection: str | None = None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DielectricFunctionHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + DielectricFunctionHandler.from_data, + DielectricFunctionHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/dielectric_tensor.py b/src/py4vasp/_calculation/dielectric_tensor.py index ef82c570d..1fc4c70a7 100644 --- a/src/py4vasp/_calculation/dielectric_tensor.py +++ b/src/py4vasp/_calculation/dielectric_tensor.py @@ -1,287 +1,287 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from typing import Optional - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation.cell import CellHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import DielectricTensor_DB -from py4vasp._util import check, convert, error -from py4vasp._util.tensor import symmetry_reduce - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - AttributeError, - TypeError, - ValueError, -) - - -class DielectricTensorHandler: - """Handler for the dielectric tensor quantity. Works with exactly one raw.DielectricTensor object.""" - - def __init__(self, raw_dielectric_tensor: raw.DielectricTensor): - self._raw_dielectric_tensor = raw_dielectric_tensor - - @classmethod - def from_data( - cls, raw_dielectric_tensor: raw.DielectricTensor - ) -> "DielectricTensorHandler": - return cls(raw_dielectric_tensor) - - def to_dict(self) -> dict: - """Read the dielectric tensor into a dictionary. - - Returns - ------- - dict - Contains the dielectric tensor and a string describing the method it - was obtained. - """ - return { - "clamped_ion": self._raw_dielectric_tensor.electron[:], - "relaxed_ion": self._read_relaxed_ion(), - "independent_particle": self._read_independent_particle(), - "method": convert.text_to_string(self._raw_dielectric_tensor.method), - } - - def __str__(self) -> str: - data = self.to_dict() - return f""" -Macroscopic static dielectric tensor (dimensionless) - {_description(data["method"])} ------------------------------------------------------- -{_dielectric_tensor_string(data["clamped_ion"], "clamped-ion")} -{_dielectric_tensor_string(data["relaxed_ion"], "relaxed-ion")} -""".strip() - - def to_database(self) -> dict: - encountered_errors = {} - error_key = "dielectric_tensor:default" - - tensor_reduced = [None, None, None] - isotropic_constant = [None, None, None] - polarizability_2d = [None, None, None] - - total_tensor, ionic_tensor, electronic_tensor = None, None, None - if not check.is_none(self._raw_dielectric_tensor.electron): - electronic_tensor = self._raw_dielectric_tensor.electron[:] - if not check.is_none(self._raw_dielectric_tensor.ion): - ionic_tensor = self._raw_dielectric_tensor.ion[:] - if not check.is_none(self._raw_dielectric_tensor.ion) and not check.is_none( - self._raw_dielectric_tensor.electron - ): - total_tensor = ( - self._raw_dielectric_tensor.electron[:] - + self._raw_dielectric_tensor.ion[:] - ) - - for idt, tensor in enumerate([total_tensor, ionic_tensor, electronic_tensor]): - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context=f"to_database.tensor[{idt}]", - ): - tensor_reduced[idt] = list(symmetry_reduce(tensor.T)) - ( - isotropic_constant[idt], - polarizability_2d[idt], - ) = self._calculate_dielectric_quantities( - tensor, - encountered_errors=encountered_errors, - error_key=error_key, - ) - - method = ( - convert.text_to_string(self._raw_dielectric_tensor.method) - if not check.is_none(self._raw_dielectric_tensor.method) - else None - ) - - return DielectricTensor_DB( - method=method, - total_3d_tensor=tensor_reduced[0], - total_3d_isotropic_dielectric_constant=isotropic_constant[0], - total_2d_polarizability=polarizability_2d[0], - ionic_3d_tensor=tensor_reduced[1], - ionic_3d_isotropic_dielectric_constant=isotropic_constant[1], - ionic_2d_polarizability=polarizability_2d[1], - electronic_3d_tensor=tensor_reduced[2], - electronic_3d_isotropic_dielectric_constant=isotropic_constant[2], - electronic_2d_polarizability=polarizability_2d[2], - ) - - # --- Private helpers --- - - def _read_relaxed_ion(self): - if check.is_none(self._raw_dielectric_tensor.ion): - return None - else: - return ( - self._raw_dielectric_tensor.electron[:] - + self._raw_dielectric_tensor.ion[:] - ) - - def _read_independent_particle(self): - if check.is_none(self._raw_dielectric_tensor.independent_particle): - return None - else: - return self._raw_dielectric_tensor.independent_particle[:] - - def _calculate_dielectric_quantities( - self, - tensor: np.ndarray, - *, - encountered_errors: Optional[dict[str, list[str]]] = None, - error_key: Optional[str] = None, - ) -> tuple: - polarizability_2d = None - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="calculate_dielectric_quantities", - ): - if not (check.is_none(self._raw_dielectric_tensor.cell)): - final_cell = CellHandler.from_data( - self._raw_dielectric_tensor.cell, steps=-1 - ) - if final_cell: - polarizability_2d = _calculate_2d_polarizability( - tensor, - final_cell, - encountered_errors=encountered_errors, - error_key=error_key, - ) - - isotropic_dielectric_constant = float(np.mean(np.diag(tensor))) - return isotropic_dielectric_constant, polarizability_2d - - -@quantity("dielectric_tensor") -class DielectricTensor: - """The dielectric tensor is the static limit of the :attr:`dielectric function`. - - The dielectric tensor represents how a material's response to an external electric - field varies with direction. It is a symmetric 3x3 matrix, encapsulating the - anisotropic nature of a material's dielectric properties. Each element of the - tensor corresponds to the dielectric function along a specific crystallographic - axis. - """ - - def __init__(self, source, quantity_name: str = "dielectric_tensor"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data( - cls, raw_dielectric_tensor: raw.DielectricTensor - ) -> "DielectricTensor": - """Create a DielectricTensor dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_dielectric_tensor)) - - def _handler_factory(self, raw_data): - return DielectricTensorHandler.from_data(raw_data) - - def read(self) -> dict: - """Read the dielectric tensor into a dictionary. - - Returns - ------- - dict - Contains the dielectric tensor and a string describing the method it - was obtained. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - DielectricTensorHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DielectricTensorHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - DielectricTensorHandler.from_data, - DielectricTensorHandler.to_database, - ) - - -def _dielectric_tensor_string(tensor, label): - if tensor is None: - return "" - row_to_string = lambda row: 6 * " " + " ".join(f"{x:12.6f}" for x in row) - rows = (row_to_string(row) for row in tensor) - return f"{label:^55}".rstrip() + "\n" + "\n".join(rows) - - -def _description(method): - if method == "dft": - return "including local field effects in DFT" - elif method == "rpa": - return "including local field effects in RPA (Hartree)" - elif method == "scf": - return "including local field effects" - elif method == "nscf": - return "excluding local field effects" - message = f"The method {method} is not implemented in this version of py4vasp." - raise exception.NotImplemented(message) - - -def _calculate_2d_polarizability( - dielectric_tensor: np.ndarray, - cell_: CellHandler, - *, - encountered_errors: Optional[dict[str, list[str]]] = None, - error_key: Optional[str] = None, -) -> float: - """ - Compute 2D polarizability (alpha_2D) for a slab system with unknown vacuum direction. - """ - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="calculate_2d_polarizability", - ): - vacuum_dir = cell_._find_likely_vacuum_direction() - if vacuum_dir is None: - return None - - eps_parallel = np.mean( - [dielectric_tensor[i, i] for i in range(3) if i != vacuum_dir] - ) - l_vacuum = np.linalg.norm(cell_.lattice_vectors()[vacuum_dir]) - - alpha_2d = (l_vacuum / (4.0 * np.pi)) * (eps_parallel - 1.0) - return alpha_2d - return None +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from typing import Optional + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation.cell import CellHandler +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import DielectricTensor_DB +from py4vasp._util import check, convert, error +from py4vasp._util.tensor import symmetry_reduce + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + AttributeError, + TypeError, + ValueError, +) + + +class DielectricTensorHandler: + """Handler for the dielectric tensor quantity. Works with exactly one raw.DielectricTensor object.""" + + def __init__(self, raw_dielectric_tensor: raw.DielectricTensor): + self._raw_dielectric_tensor = raw_dielectric_tensor + + @classmethod + def from_data( + cls, raw_dielectric_tensor: raw.DielectricTensor + ) -> "DielectricTensorHandler": + return cls(raw_dielectric_tensor) + + def to_dict(self) -> dict: + """Read the dielectric tensor into a dictionary. + + Returns + ------- + dict + Contains the dielectric tensor and a string describing the method it + was obtained. + """ + return { + "clamped_ion": self._raw_dielectric_tensor.electron[:], + "relaxed_ion": self._read_relaxed_ion(), + "independent_particle": self._read_independent_particle(), + "method": convert.text_to_string(self._raw_dielectric_tensor.method), + } + + def __str__(self) -> str: + data = self.to_dict() + return f""" +Macroscopic static dielectric tensor (dimensionless) + {_description(data["method"])} +------------------------------------------------------ +{_dielectric_tensor_string(data["clamped_ion"], "clamped-ion")} +{_dielectric_tensor_string(data["relaxed_ion"], "relaxed-ion")} +""".strip() + + def to_database(self) -> dict: + encountered_errors = {} + error_key = "dielectric_tensor:default" + + tensor_reduced = [None, None, None] + isotropic_constant = [None, None, None] + polarizability_2d = [None, None, None] + + total_tensor, ionic_tensor, electronic_tensor = None, None, None + if not check.is_none(self._raw_dielectric_tensor.electron): + electronic_tensor = self._raw_dielectric_tensor.electron[:] + if not check.is_none(self._raw_dielectric_tensor.ion): + ionic_tensor = self._raw_dielectric_tensor.ion[:] + if not check.is_none(self._raw_dielectric_tensor.ion) and not check.is_none( + self._raw_dielectric_tensor.electron + ): + total_tensor = ( + self._raw_dielectric_tensor.electron[:] + + self._raw_dielectric_tensor.ion[:] + ) + + for idt, tensor in enumerate([total_tensor, ionic_tensor, electronic_tensor]): + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context=f"to_database.tensor[{idt}]", + ): + tensor_reduced[idt] = list(symmetry_reduce(tensor.T)) + ( + isotropic_constant[idt], + polarizability_2d[idt], + ) = self._calculate_dielectric_quantities( + tensor, + encountered_errors=encountered_errors, + error_key=error_key, + ) + + method = ( + convert.text_to_string(self._raw_dielectric_tensor.method) + if not check.is_none(self._raw_dielectric_tensor.method) + else None + ) + + return DielectricTensor_DB( + method=method, + total_3d_tensor=tensor_reduced[0], + total_3d_isotropic_dielectric_constant=isotropic_constant[0], + total_2d_polarizability=polarizability_2d[0], + ionic_3d_tensor=tensor_reduced[1], + ionic_3d_isotropic_dielectric_constant=isotropic_constant[1], + ionic_2d_polarizability=polarizability_2d[1], + electronic_3d_tensor=tensor_reduced[2], + electronic_3d_isotropic_dielectric_constant=isotropic_constant[2], + electronic_2d_polarizability=polarizability_2d[2], + ) + + # --- Private helpers --- + + def _read_relaxed_ion(self): + if check.is_none(self._raw_dielectric_tensor.ion): + return None + else: + return ( + self._raw_dielectric_tensor.electron[:] + + self._raw_dielectric_tensor.ion[:] + ) + + def _read_independent_particle(self): + if check.is_none(self._raw_dielectric_tensor.independent_particle): + return None + else: + return self._raw_dielectric_tensor.independent_particle[:] + + def _calculate_dielectric_quantities( + self, + tensor: np.ndarray, + *, + encountered_errors: Optional[dict[str, list[str]]] = None, + error_key: Optional[str] = None, + ) -> tuple: + polarizability_2d = None + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="calculate_dielectric_quantities", + ): + if not (check.is_none(self._raw_dielectric_tensor.cell)): + final_cell = CellHandler.from_data( + self._raw_dielectric_tensor.cell, steps=-1 + ) + if final_cell: + polarizability_2d = _calculate_2d_polarizability( + tensor, + final_cell, + encountered_errors=encountered_errors, + error_key=error_key, + ) + + isotropic_dielectric_constant = float(np.mean(np.diag(tensor))) + return isotropic_dielectric_constant, polarizability_2d + + +@quantity("dielectric_tensor") +class DielectricTensor: + """The dielectric tensor is the static limit of the :attr:`dielectric function`. + + The dielectric tensor represents how a material's response to an external electric + field varies with direction. It is a symmetric 3x3 matrix, encapsulating the + anisotropic nature of a material's dielectric properties. Each element of the + tensor corresponds to the dielectric function along a specific crystallographic + axis. + """ + + def __init__(self, source, quantity_name: str = "dielectric_tensor"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data( + cls, raw_dielectric_tensor: raw.DielectricTensor + ) -> "DielectricTensor": + """Create a DielectricTensor dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_dielectric_tensor)) + + def _handler_factory(self, raw_data): + return DielectricTensorHandler.from_data(raw_data) + + def read(self) -> dict: + """Read the dielectric tensor into a dictionary. + + Returns + ------- + dict + Contains the dielectric tensor and a string describing the method it + was obtained. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + DielectricTensorHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DielectricTensorHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + DielectricTensorHandler.from_data, + DielectricTensorHandler.to_database, + ) + + +def _dielectric_tensor_string(tensor, label): + if tensor is None: + return "" + row_to_string = lambda row: 6 * " " + " ".join(f"{x:12.6f}" for x in row) + rows = (row_to_string(row) for row in tensor) + return f"{label:^55}".rstrip() + "\n" + "\n".join(rows) + + +def _description(method): + if method == "dft": + return "including local field effects in DFT" + elif method == "rpa": + return "including local field effects in RPA (Hartree)" + elif method == "scf": + return "including local field effects" + elif method == "nscf": + return "excluding local field effects" + message = f"The method {method} is not implemented in this version of py4vasp." + raise exception.NotImplemented(message) + + +def _calculate_2d_polarizability( + dielectric_tensor: np.ndarray, + cell_: CellHandler, + *, + encountered_errors: Optional[dict[str, list[str]]] = None, + error_key: Optional[str] = None, +) -> float: + """ + Compute 2D polarizability (alpha_2D) for a slab system with unknown vacuum direction. + """ + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="calculate_2d_polarizability", + ): + vacuum_dir = cell_._find_likely_vacuum_direction() + if vacuum_dir is None: + return None + + eps_parallel = np.mean( + [dielectric_tensor[i, i] for i in range(3) if i != vacuum_dir] + ) + l_vacuum = np.linalg.norm(cell_.lattice_vectors()[vacuum_dir]) + + alpha_2d = (l_vacuum / (4.0 * np.pi)) * (eps_parallel - 1.0) + return alpha_2d + return None diff --git a/src/py4vasp/_calculation/dos.py b/src/py4vasp/_calculation/dos.py index 9856321d9..fceab3a1f 100644 --- a/src/py4vasp/_calculation/dos.py +++ b/src/py4vasp/_calculation/dos.py @@ -1,582 +1,582 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import pathlib -from contextlib import suppress - -import numpy as np - -from py4vasp import exception -from py4vasp._calculation import projector -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.projector import ProjectorHandler -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Dos_DB -from py4vasp._third_party import graph -from py4vasp._util import check, documentation, import_ - -pd = import_.optional("pandas") -pretty = import_.optional("IPython.lib.pretty") - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - AttributeError, - TypeError, - ValueError, - IndexError, - ZeroDivisionError, -) - - -class DosHandler: - """Handler for density of states data.""" - - def __init__(self, raw_dos: raw.Dos): - self._raw_dos = raw_dos - - @classmethod - def from_data(cls, raw_dos: raw.Dos) -> "DosHandler": - return cls(raw_dos) - - def __str__(self): - energies = self._raw_dos.energies - if self._is_collinear(): - label = "collinear Dos" - elif self._is_noncollinear(): - label = "noncollinear Dos" - else: - label = "Dos" - return f"""\ -{label}: - energies: [{energies[0]:0.2f}, {energies[-1]:0.2f}] {len(energies)} points -{str(self._projector())}""" - - def to_dict(self, selection=None) -> dict: - data = self._read_data(selection) - data.pop(projector.SPIN_PROJECTION, None) - return {**data, "fermi_energy": self._raw_dos.fermi_energy} - - def to_database(self, fermi_energy=None) -> dict: - raw_fermi_energy = ( - self._raw_dos.fermi_energy - if not check.is_none(self._raw_dos.fermi_energy) - else None - ) - dos_at_fermi_dict = self._dos_at_energy(fermi_energy or raw_fermi_energy) - dos_at_raw_fermi_dict = self._dos_at_energy(raw_fermi_energy) - - dos_at_fermi_total = dos_at_fermi_dict.get("total", None) - dos_at_raw_fermi_total = dos_at_raw_fermi_dict.get("total", None) - dos_at_fermi_up = dos_at_fermi_dict.get("up", None) - dos_at_fermi_down = dos_at_fermi_dict.get("down", None) - dos_at_raw_fermi_up = dos_at_raw_fermi_dict.get("up", None) - dos_at_raw_fermi_down = dos_at_raw_fermi_dict.get("down", None) - - return Dos_DB( - dos_at_fermi_total=dos_at_fermi_total, - dos_at_fermi_up=dos_at_fermi_up, - dos_at_fermi_down=dos_at_fermi_down, - dos_at_raw_fermi_total=dos_at_raw_fermi_total, - dos_at_raw_fermi_up=dos_at_raw_fermi_up, - dos_at_raw_fermi_down=dos_at_raw_fermi_down, - energy_min=( - float(np.min(self._raw_dos.energies[:])) - if not check.is_none(self._raw_dos.energies) - else None - ), - energy_max=( - float(np.max(self._raw_dos.energies[:])) - if not check.is_none(self._raw_dos.energies) - else None - ), - ) - - def to_graph(self, selection=None) -> graph.Graph: - data = self._read_data(selection) - energies = data.pop("energies") - data.pop(projector.SPIN_PROJECTION, None) - return graph.Graph( - series=list(_series(energies, data)), - xlabel="Energy (eV)", - ylabel="DOS (1/eV)", - ) - - def to_frame(self, selection=None): - data = self._read_data(selection) - data.pop(projector.SPIN_PROJECTION, None) - df = pd.DataFrame(data) - df.fermi_energy = self._raw_dos.fermi_energy - return df - - def selections(self) -> dict: - return self._projector().selections() - - def _is_collinear(self): - return len(self._raw_dos.dos) == 2 - - def _is_noncollinear(self): - return len(self._raw_dos.dos) == 4 - - def _projector(self): - return ProjectorHandler.from_data(self._raw_dos.projectors) - - def _read_data(self, selection): - return { - **self._read_energies(), - **self._read_total_dos(), - **self._projector().project(selection, self._raw_dos.projections), - } - - def _read_energies(self): - return {"energies": self._raw_dos.energies[:] - self._raw_dos.fermi_energy} - - def _read_total_dos(self): - if self._is_collinear(): - return {"up": self._raw_dos.dos[0, :], "down": self._raw_dos.dos[1, :]} - else: - return {"total": self._raw_dos.dos[0, :]} - - def _dos_at_energy(self, energy): - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - energies = self._raw_dos.energies[:] - dos_dict = self._read_total_dos() - dos_at_energy = {} - for key, dos in dos_dict.items(): - idx = (np.abs(energies - energy)).argmin() - if energies[idx] == energy: - dos_at_energy[key] = float(dos[idx]) - else: - if energies[idx] < energy: - idx_low = idx - idx_high = idx + 1 - else: - idx_low = idx - 1 - idx_high = idx - if (idx_low < 0) or (idx_high >= len(energies)): - dos_at_energy[key] = None - continue - dos_low = dos[idx_low] - dos_high = dos[idx_high] - energy_low = energies[idx_low] - energy_high = energies[idx_high] - dos_at_energy[key] = float( - dos_low - + (dos_high - dos_low) - * (energy - energy_low) - / (energy_high - energy_low) - ) - return dos_at_energy - return {} - - -@quantity("dos") -class Dos(graph.Mixin): - """The density of states (DOS) describes the number of states per energy. - - The DOS quantifies the distribution of electronic states within an energy range - in a material. It provides information about the number of electronic states at - each energy level and offers insights into the material's electronic - structure. On-site projections near the atoms (projected DOS) offer a more detailed - view. This analysis breaks down the DOS contributions by atom, orbital and spin. - Investigating the projected DOS is often a useful step to understand the - electronic properties because it shows how different orbitals and elements - contribute and influence the material's properties. - - VASP writes the DOS after every calculation and the projected DOS if you set - :tag:`LORBIT` in the INCAR file. You can use this class to extract this data. - Typically you want to run a non self consistent calculation with a denser - mesh for a smoother DOS but the class will work independent of it. If you - generated a projected DOS, you can use this class to select which subset of - these orbitals to read or plot. - - Examples - -------- - - First, we create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP with :tag:`LORBIT`. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you want to visualize the total DOS, you can use the `plot` method. This will - show the different spin components if :tag:`ISPIN` = 2 - - >>> calculation.dos.plot() - Graph(series=[Series(x=array(...), y=array(...), label='total', ...)], - xlabel='Energy (eV)', ..., ylabel='DOS (1/eV)', ...) - - If you need the raw data, you can read the DOS into a Python dictionary - - >>> calculation.dos.read() - {'energies': array(...), 'total': array(...), 'fermi_energy': ...} - - These methods also accept selections for specific orbitals if you used VASP with - :tag:`LORBIT`. You can get a list of the allowed choices with - - >>> calculation.dos.selections() - {'dos': ['default', 'kpoints_opt'], 'atom': [...], 'orbital': [...], 'spin': [...]} - """ - - def __init__(self, source, quantity_name="dos"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_dos): - return cls(source=DataSource(raw_dos)) - - def _handler_factory(self, raw): - return DosHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DosHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - @documentation.format(selection_doc=projector.selection_doc) - def read(self, selection=None) -> dict: - """Read the DOS into a dictionary. - - You will always get an "energies" component that describes the energy mesh for - the density of states. The energies are shifted with respect to VASP such that - the Fermi energy is at 0. py4vasp returns also the original "fermi_energy" so - you can revert this if you want. If :tag:`ISPIN` = 2, you will get the total - DOS spin resolved as "up" and "down" component. Otherwise, you will get just - the "total" DOS. When you set :tag:`LORBIT` in the INCAR file and pass in a - selection, you will obtain the projected DOS with a label corresponding to the - projection. - - We create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP with :tag:`LORBIT`. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> collinear_calculation = demo.calculation(path, selection="collinear") - >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") - - Parameters - ---------- - {selection_doc} - - Returns - ------- - dict - Contains the energies at which the DOS was evaluated aligned to the - Fermi energy and the total DOS or the spin-resolved DOS for - spin-polarized calculations. If available and a selection is passed, - the orbital resolved DOS for the selected orbitals is included. - - Examples - -------- - - To obtain the total DOS along with the energy mesh and the Fermi energy you - do not need any arguments. For :tag:`ISPIN` = 2, this will "up" and "down" - DOS as two separate entries. - - >>> calculation.dos.read() - {{'energies': array(...), 'total': array(...), 'fermi_energy': ...}} - - Select the p orbitals of the first atom in the POSCAR file: - - >>> calculation.dos.read(selection="1(p)") - {{'energies': array(...), 'total': array(...), 'Sr_1_p': array(...), - 'fermi_energy': ...}} - - Select the d orbitals of Sr and Ti: - - >>> calculation.dos.read("d(Sr, Ti)") - {{'energies': array(...), 'total': array(...), 'Sr_d': array(...), - 'Ti_d': array(...), 'fermi_energy': ...}} - - Collinear calculations separate the DOS into two spin components - - >>> collinear_calculation.dos.read() - {{'energies': array(...), 'up': array(...), 'down': array(...), ...}} - - You can also select spin contribution of specific atoms or orbitals, e.g. the - spin-up contribution of the first three atoms combined - - >>> collinear_calculation.dos.read("up(1:3)") - {{'energies': array(...), 'up': array(...), 'down': array(...), - '1:3_up': array(...), 'fermi_energy': ...}} - - Noncollinear calculations contain four spin components but by default only - the total DOS is shown - - >>> noncollinear_calculation.dos.read() - {{'energies': array(...), 'total': array(...), 'fermi_energy': ...}} - - You can select specific spin components if you want, e.g. the spin projection - along the z axis - - >>> noncollinear_calculation.dos.read("sigma_z") - {{'energies': array(...), 'total': array(...), 'sigma_z': array(...), - 'fermi_energy': ...}} - - You can also use simple addition and subtraction to combine the contributions of - different orbitals, e.g. add the contribution of three d orbitals - - >>> calculation.dos.read("dxy + dxz + dyz") - {{'energies': array(...), 'total': array(...), 'dxy + dxz + dyz': array(...), - 'fermi_energy': ...}} - - Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT - file (analogously for KPOINTS_WAN) - - >>> calculation.dos.read("kpoints_opt") - {{'energies': array(...), 'total': array(...), 'fermi_energy': ...}} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DosHandler.to_dict, - ) - - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - @documentation.format(selection_doc=projector.selection_doc) - def to_graph(self, selection=None) -> graph.Graph: - """Read the DOS and convert it into a graph. - - The x axis is the energy mesh used in the calculation shifted such that the - Fermi energy is at 0. On the y axis, we show the DOS. For :tag:`ISPIN` = 2, the - different spin components are shown with opposite sign: "up" with a positive - sign and "down" with a negative one. If you used :tag:`LORBIT` in your VASP - calculation and you pass in a selection, py4vasp will add additional lines - corresponding to the selected projections. - - We create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP with :tag:`LORBIT`. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> collinear_calculation = demo.calculation(path, selection="collinear") - >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") - - Parameters - ---------- - {selection_doc} - - Returns - ------- - Graph - Graph containing the total DOS. If the calculation was spin polarized, - the resulting DOS is spin resolved and the spin-down DOS is plotted - towards negative values. If a selection is given the orbital-resolved - DOS is given for the specified projectors. - - Examples - -------- - - For the total DOS, you do not need any arguments. py4vasp will automatically - use two separate lines, if you used :tag:`ISPIN` = 2 in the VASP calculation - - >>> calculation.dos.to_graph() - Graph(series=[Series(x=array(...), y=array(...), label='total', ...)], - xlabel='Energy (eV)', ..., ylabel='DOS (1/eV)', ...) - - Select the p orbitals of the first atom in the POSCAR file: - - >>> calculation.dos.to_graph(selection="1(p)") - Graph(series=[Series(..., label='total', ...), Series(..., label='Sr_1_p', ...)], ...) - - Select the d orbitals of Sr and Ti: - - >>> calculation.dos.to_graph("d(Sr, Ti)") - Graph(series=[Series(..., label='total', ...), Series(..., label='Sr_d', ...), - Series(..., label='Ti_d', ...)], ...) - - Collinear calculations separate the DOS into two spin components - - >>> collinear_calculation.dos.to_graph() - Graph(series=[Series(..., label='up', ...), Series(..., label='down', ...)], ...) - - You can also select spin contribution of specific atoms or orbitals, e.g. the - spin-up contribution of the first three atoms combined - - >>> collinear_calculation.dos.to_graph("up(1:3)") - Graph(series=[Series(..., label='up', ...), Series(..., label='down', ...), - Series(..., label='1:3_up', ...)], ...) - - Noncollinear calculations contain four spin components but by default only - the total DOS is shown - - >>> noncollinear_calculation.dos.to_graph() - Graph(series=[Series(..., label='total', ...)], ...) - - You can select specific spin components if you want, e.g. the spin projection - along the z axis - - >>> noncollinear_calculation.dos.to_graph("sigma_z") - Graph(series=[Series(..., label='total', ...), Series(..., label='sigma_z', ...)], - ...) - - You can also use simple addition and subtraction to combine the contributions of - different orbitals, e.g. add the contribution of three d orbitals - - >>> calculation.dos.to_graph("dxy + dxz + dyz") - Graph(series=[Series(..., label='total', ...), Series(..., label='dxy + dxz + dyz', - ...)], ...) - - Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT - file (analogously for KPOINTS_WAN) - - >>> calculation.dos.to_graph("kpoints_opt") - Graph(series=[Series(..., label='total', ...)], ...) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DosHandler.to_graph, - ) - - @documentation.format(selection_doc=projector.selection_doc) - def to_frame(self, selection=None): - """Read the data into a pandas DataFrame. - - We create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP with :tag:`LORBIT`. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> collinear_calculation = demo.calculation(path, selection="collinear") - >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") - - Parameters - ---------- - {selection_doc} - - Returns - ------- - pd.DataFrame - Contains the energies at which the DOS was evaluated aligned to the - Fermi energy and the total DOS or the spin-resolved DOS for - spin-polarized calculations. If available and a selection is passed, - the orbital resolved DOS for the selected orbitals is included. - - Examples - -------- - - >>> calculation.dos.to_frame() - energies total - 0 ... - - Select the p orbitals of the first atom in the POSCAR file: - - >>> calculation.dos.to_frame(selection="1(p)") - energies total Sr_1_p - 0 ... - - Select the d orbitals of Sr and Ti: - - >>> calculation.dos.to_frame("d(Sr, Ti)") - energies total Sr_d Ti_d - 0 ... - - Collinear calculations separate the DOS into two spin components - - >>> collinear_calculation.dos.to_frame() - energies up down - 0 ... - - You can also select spin contribution of specific atoms or orbitals, e.g. the - spin-up contribution of the first three atoms combined - - >>> collinear_calculation.dos.to_frame("up(1:3)") - energies up down 1:3_up - 0 ... - - Noncollinear calculations contain four spin components but by default only - the total DOS is shown - - >>> noncollinear_calculation.dos.to_frame() - energies total - 0 ... - - You can select specific spin components if you want, e.g. the spin projection - along the z axis - - >>> noncollinear_calculation.dos.to_frame("sigma_z") - energies total sigma_z - 0 ... - - You can also use simple addition and subtraction to combine the contributions of - different orbitals, e.g. add the contribution of three d orbitals - - >>> calculation.dos.to_frame("dxy + dxz + dyz") - energies total dxy + dxz + dyz - 0 ... - - Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT - file (analogously for KPOINTS_WAN) - - >>> calculation.dos.to_frame("kpoints_opt") - energies total - 0 ... - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DosHandler.to_frame, - ) - - def selections(self, selection=None) -> dict: - from py4vasp._raw import definition as raw_module - - handler_selections = merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - DosHandler.selections, - ) - sources = list(raw_module.selections(self._quantity_name)) - return {self._quantity_name: sources, **handler_selections} - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - DosHandler.from_data, - DosHandler.to_database, - ) - - -def _series(energies, data): - for name, dos in data.items(): - spin_factor = -1 if _flip_down_component(name) else 1 - yield graph.Series(energies, spin_factor * dos, name) - - -def _flip_down_component(name): - return "down" in name and "up" not in name and "total" not in name +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import pathlib +from contextlib import suppress + +import numpy as np + +from py4vasp import exception +from py4vasp._calculation import projector +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.projector import ProjectorHandler +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Dos_DB +from py4vasp._third_party import graph +from py4vasp._util import check, documentation, import_ + +pd = import_.optional("pandas") +pretty = import_.optional("IPython.lib.pretty") + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + AttributeError, + TypeError, + ValueError, + IndexError, + ZeroDivisionError, +) + + +class DosHandler: + """Handler for density of states data.""" + + def __init__(self, raw_dos: raw.Dos): + self._raw_dos = raw_dos + + @classmethod + def from_data(cls, raw_dos: raw.Dos) -> "DosHandler": + return cls(raw_dos) + + def __str__(self): + energies = self._raw_dos.energies + if self._is_collinear(): + label = "collinear Dos" + elif self._is_noncollinear(): + label = "noncollinear Dos" + else: + label = "Dos" + return f"""\ +{label}: + energies: [{energies[0]:0.2f}, {energies[-1]:0.2f}] {len(energies)} points +{str(self._projector())}""" + + def to_dict(self, selection=None) -> dict: + data = self._read_data(selection) + data.pop(projector.SPIN_PROJECTION, None) + return {**data, "fermi_energy": self._raw_dos.fermi_energy} + + def to_database(self, fermi_energy=None) -> dict: + raw_fermi_energy = ( + self._raw_dos.fermi_energy + if not check.is_none(self._raw_dos.fermi_energy) + else None + ) + dos_at_fermi_dict = self._dos_at_energy(fermi_energy or raw_fermi_energy) + dos_at_raw_fermi_dict = self._dos_at_energy(raw_fermi_energy) + + dos_at_fermi_total = dos_at_fermi_dict.get("total", None) + dos_at_raw_fermi_total = dos_at_raw_fermi_dict.get("total", None) + dos_at_fermi_up = dos_at_fermi_dict.get("up", None) + dos_at_fermi_down = dos_at_fermi_dict.get("down", None) + dos_at_raw_fermi_up = dos_at_raw_fermi_dict.get("up", None) + dos_at_raw_fermi_down = dos_at_raw_fermi_dict.get("down", None) + + return Dos_DB( + dos_at_fermi_total=dos_at_fermi_total, + dos_at_fermi_up=dos_at_fermi_up, + dos_at_fermi_down=dos_at_fermi_down, + dos_at_raw_fermi_total=dos_at_raw_fermi_total, + dos_at_raw_fermi_up=dos_at_raw_fermi_up, + dos_at_raw_fermi_down=dos_at_raw_fermi_down, + energy_min=( + float(np.min(self._raw_dos.energies[:])) + if not check.is_none(self._raw_dos.energies) + else None + ), + energy_max=( + float(np.max(self._raw_dos.energies[:])) + if not check.is_none(self._raw_dos.energies) + else None + ), + ) + + def to_graph(self, selection=None) -> graph.Graph: + data = self._read_data(selection) + energies = data.pop("energies") + data.pop(projector.SPIN_PROJECTION, None) + return graph.Graph( + series=list(_series(energies, data)), + xlabel="Energy (eV)", + ylabel="DOS (1/eV)", + ) + + def to_frame(self, selection=None): + data = self._read_data(selection) + data.pop(projector.SPIN_PROJECTION, None) + df = pd.DataFrame(data) + df.fermi_energy = self._raw_dos.fermi_energy + return df + + def selections(self) -> dict: + return self._projector().selections() + + def _is_collinear(self): + return len(self._raw_dos.dos) == 2 + + def _is_noncollinear(self): + return len(self._raw_dos.dos) == 4 + + def _projector(self): + return ProjectorHandler.from_data(self._raw_dos.projectors) + + def _read_data(self, selection): + return { + **self._read_energies(), + **self._read_total_dos(), + **self._projector().project(selection, self._raw_dos.projections), + } + + def _read_energies(self): + return {"energies": self._raw_dos.energies[:] - self._raw_dos.fermi_energy} + + def _read_total_dos(self): + if self._is_collinear(): + return {"up": self._raw_dos.dos[0, :], "down": self._raw_dos.dos[1, :]} + else: + return {"total": self._raw_dos.dos[0, :]} + + def _dos_at_energy(self, energy): + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + energies = self._raw_dos.energies[:] + dos_dict = self._read_total_dos() + dos_at_energy = {} + for key, dos in dos_dict.items(): + idx = (np.abs(energies - energy)).argmin() + if energies[idx] == energy: + dos_at_energy[key] = float(dos[idx]) + else: + if energies[idx] < energy: + idx_low = idx + idx_high = idx + 1 + else: + idx_low = idx - 1 + idx_high = idx + if (idx_low < 0) or (idx_high >= len(energies)): + dos_at_energy[key] = None + continue + dos_low = dos[idx_low] + dos_high = dos[idx_high] + energy_low = energies[idx_low] + energy_high = energies[idx_high] + dos_at_energy[key] = float( + dos_low + + (dos_high - dos_low) + * (energy - energy_low) + / (energy_high - energy_low) + ) + return dos_at_energy + return {} + + +@quantity("dos") +class Dos(graph.Mixin): + """The density of states (DOS) describes the number of states per energy. + + The DOS quantifies the distribution of electronic states within an energy range + in a material. It provides information about the number of electronic states at + each energy level and offers insights into the material's electronic + structure. On-site projections near the atoms (projected DOS) offer a more detailed + view. This analysis breaks down the DOS contributions by atom, orbital and spin. + Investigating the projected DOS is often a useful step to understand the + electronic properties because it shows how different orbitals and elements + contribute and influence the material's properties. + + VASP writes the DOS after every calculation and the projected DOS if you set + :tag:`LORBIT` in the INCAR file. You can use this class to extract this data. + Typically you want to run a non self consistent calculation with a denser + mesh for a smoother DOS but the class will work independent of it. If you + generated a projected DOS, you can use this class to select which subset of + these orbitals to read or plot. + + Examples + -------- + + First, we create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP with :tag:`LORBIT`. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you want to visualize the total DOS, you can use the `plot` method. This will + show the different spin components if :tag:`ISPIN` = 2 + + >>> calculation.dos.plot() + Graph(series=[Series(x=array(...), y=array(...), label='total', ...)], + xlabel='Energy (eV)', ..., ylabel='DOS (1/eV)', ...) + + If you need the raw data, you can read the DOS into a Python dictionary + + >>> calculation.dos.read() + {'energies': array(...), 'total': array(...), 'fermi_energy': ...} + + These methods also accept selections for specific orbitals if you used VASP with + :tag:`LORBIT`. You can get a list of the allowed choices with + + >>> calculation.dos.selections() + {'dos': ['default', 'kpoints_opt'], 'atom': [...], 'orbital': [...], 'spin': [...]} + """ + + def __init__(self, source, quantity_name="dos"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_dos): + return cls(source=DataSource(raw_dos)) + + def _handler_factory(self, raw): + return DosHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DosHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + @documentation.format(selection_doc=projector.selection_doc) + def read(self, selection=None) -> dict: + """Read the DOS into a dictionary. + + You will always get an "energies" component that describes the energy mesh for + the density of states. The energies are shifted with respect to VASP such that + the Fermi energy is at 0. py4vasp returns also the original "fermi_energy" so + you can revert this if you want. If :tag:`ISPIN` = 2, you will get the total + DOS spin resolved as "up" and "down" component. Otherwise, you will get just + the "total" DOS. When you set :tag:`LORBIT` in the INCAR file and pass in a + selection, you will obtain the projected DOS with a label corresponding to the + projection. + + We create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP with :tag:`LORBIT`. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> collinear_calculation = demo.calculation(path, selection="collinear") + >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") + + Parameters + ---------- + {selection_doc} + + Returns + ------- + dict + Contains the energies at which the DOS was evaluated aligned to the + Fermi energy and the total DOS or the spin-resolved DOS for + spin-polarized calculations. If available and a selection is passed, + the orbital resolved DOS for the selected orbitals is included. + + Examples + -------- + + To obtain the total DOS along with the energy mesh and the Fermi energy you + do not need any arguments. For :tag:`ISPIN` = 2, this will "up" and "down" + DOS as two separate entries. + + >>> calculation.dos.read() + {{'energies': array(...), 'total': array(...), 'fermi_energy': ...}} + + Select the p orbitals of the first atom in the POSCAR file: + + >>> calculation.dos.read(selection="1(p)") + {{'energies': array(...), 'total': array(...), 'Sr_1_p': array(...), + 'fermi_energy': ...}} + + Select the d orbitals of Sr and Ti: + + >>> calculation.dos.read("d(Sr, Ti)") + {{'energies': array(...), 'total': array(...), 'Sr_d': array(...), + 'Ti_d': array(...), 'fermi_energy': ...}} + + Collinear calculations separate the DOS into two spin components + + >>> collinear_calculation.dos.read() + {{'energies': array(...), 'up': array(...), 'down': array(...), ...}} + + You can also select spin contribution of specific atoms or orbitals, e.g. the + spin-up contribution of the first three atoms combined + + >>> collinear_calculation.dos.read("up(1:3)") + {{'energies': array(...), 'up': array(...), 'down': array(...), + '1:3_up': array(...), 'fermi_energy': ...}} + + Noncollinear calculations contain four spin components but by default only + the total DOS is shown + + >>> noncollinear_calculation.dos.read() + {{'energies': array(...), 'total': array(...), 'fermi_energy': ...}} + + You can select specific spin components if you want, e.g. the spin projection + along the z axis + + >>> noncollinear_calculation.dos.read("sigma_z") + {{'energies': array(...), 'total': array(...), 'sigma_z': array(...), + 'fermi_energy': ...}} + + You can also use simple addition and subtraction to combine the contributions of + different orbitals, e.g. add the contribution of three d orbitals + + >>> calculation.dos.read("dxy + dxz + dyz") + {{'energies': array(...), 'total': array(...), 'dxy + dxz + dyz': array(...), + 'fermi_energy': ...}} + + Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT + file (analogously for KPOINTS_WAN) + + >>> calculation.dos.read("kpoints_opt") + {{'energies': array(...), 'total': array(...), 'fermi_energy': ...}} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DosHandler.to_dict, + ) + + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + @documentation.format(selection_doc=projector.selection_doc) + def to_graph(self, selection=None) -> graph.Graph: + """Read the DOS and convert it into a graph. + + The x axis is the energy mesh used in the calculation shifted such that the + Fermi energy is at 0. On the y axis, we show the DOS. For :tag:`ISPIN` = 2, the + different spin components are shown with opposite sign: "up" with a positive + sign and "down" with a negative one. If you used :tag:`LORBIT` in your VASP + calculation and you pass in a selection, py4vasp will add additional lines + corresponding to the selected projections. + + We create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP with :tag:`LORBIT`. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> collinear_calculation = demo.calculation(path, selection="collinear") + >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") + + Parameters + ---------- + {selection_doc} + + Returns + ------- + Graph + Graph containing the total DOS. If the calculation was spin polarized, + the resulting DOS is spin resolved and the spin-down DOS is plotted + towards negative values. If a selection is given the orbital-resolved + DOS is given for the specified projectors. + + Examples + -------- + + For the total DOS, you do not need any arguments. py4vasp will automatically + use two separate lines, if you used :tag:`ISPIN` = 2 in the VASP calculation + + >>> calculation.dos.to_graph() + Graph(series=[Series(x=array(...), y=array(...), label='total', ...)], + xlabel='Energy (eV)', ..., ylabel='DOS (1/eV)', ...) + + Select the p orbitals of the first atom in the POSCAR file: + + >>> calculation.dos.to_graph(selection="1(p)") + Graph(series=[Series(..., label='total', ...), Series(..., label='Sr_1_p', ...)], ...) + + Select the d orbitals of Sr and Ti: + + >>> calculation.dos.to_graph("d(Sr, Ti)") + Graph(series=[Series(..., label='total', ...), Series(..., label='Sr_d', ...), + Series(..., label='Ti_d', ...)], ...) + + Collinear calculations separate the DOS into two spin components + + >>> collinear_calculation.dos.to_graph() + Graph(series=[Series(..., label='up', ...), Series(..., label='down', ...)], ...) + + You can also select spin contribution of specific atoms or orbitals, e.g. the + spin-up contribution of the first three atoms combined + + >>> collinear_calculation.dos.to_graph("up(1:3)") + Graph(series=[Series(..., label='up', ...), Series(..., label='down', ...), + Series(..., label='1:3_up', ...)], ...) + + Noncollinear calculations contain four spin components but by default only + the total DOS is shown + + >>> noncollinear_calculation.dos.to_graph() + Graph(series=[Series(..., label='total', ...)], ...) + + You can select specific spin components if you want, e.g. the spin projection + along the z axis + + >>> noncollinear_calculation.dos.to_graph("sigma_z") + Graph(series=[Series(..., label='total', ...), Series(..., label='sigma_z', ...)], + ...) + + You can also use simple addition and subtraction to combine the contributions of + different orbitals, e.g. add the contribution of three d orbitals + + >>> calculation.dos.to_graph("dxy + dxz + dyz") + Graph(series=[Series(..., label='total', ...), Series(..., label='dxy + dxz + dyz', + ...)], ...) + + Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT + file (analogously for KPOINTS_WAN) + + >>> calculation.dos.to_graph("kpoints_opt") + Graph(series=[Series(..., label='total', ...)], ...) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DosHandler.to_graph, + ) + + @documentation.format(selection_doc=projector.selection_doc) + def to_frame(self, selection=None): + """Read the data into a pandas DataFrame. + + We create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP with :tag:`LORBIT`. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> collinear_calculation = demo.calculation(path, selection="collinear") + >>> noncollinear_calculation = demo.calculation(path, selection="noncollinear") + + Parameters + ---------- + {selection_doc} + + Returns + ------- + pd.DataFrame + Contains the energies at which the DOS was evaluated aligned to the + Fermi energy and the total DOS or the spin-resolved DOS for + spin-polarized calculations. If available and a selection is passed, + the orbital resolved DOS for the selected orbitals is included. + + Examples + -------- + + >>> calculation.dos.to_frame() + energies total + 0 ... + + Select the p orbitals of the first atom in the POSCAR file: + + >>> calculation.dos.to_frame(selection="1(p)") + energies total Sr_1_p + 0 ... + + Select the d orbitals of Sr and Ti: + + >>> calculation.dos.to_frame("d(Sr, Ti)") + energies total Sr_d Ti_d + 0 ... + + Collinear calculations separate the DOS into two spin components + + >>> collinear_calculation.dos.to_frame() + energies up down + 0 ... + + You can also select spin contribution of specific atoms or orbitals, e.g. the + spin-up contribution of the first three atoms combined + + >>> collinear_calculation.dos.to_frame("up(1:3)") + energies up down 1:3_up + 0 ... + + Noncollinear calculations contain four spin components but by default only + the total DOS is shown + + >>> noncollinear_calculation.dos.to_frame() + energies total + 0 ... + + You can select specific spin components if you want, e.g. the spin projection + along the z axis + + >>> noncollinear_calculation.dos.to_frame("sigma_z") + energies total sigma_z + 0 ... + + You can also use simple addition and subtraction to combine the contributions of + different orbitals, e.g. add the contribution of three d orbitals + + >>> calculation.dos.to_frame("dxy + dxz + dyz") + energies total dxy + dxz + dyz + 0 ... + + Read the density of states generated by the '''k'''-point mesh in the KPOINTS_OPT + file (analogously for KPOINTS_WAN) + + >>> calculation.dos.to_frame("kpoints_opt") + energies total + 0 ... + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DosHandler.to_frame, + ) + + def selections(self, selection=None) -> dict: + from py4vasp._raw import definition as raw_module + + handler_selections = merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + DosHandler.selections, + ) + sources = list(raw_module.selections(self._quantity_name)) + return {self._quantity_name: sources, **handler_selections} + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + DosHandler.from_data, + DosHandler.to_database, + ) + + +def _series(energies, data): + for name, dos in data.items(): + spin_factor = -1 if _flip_down_component(name) else 1 + yield graph.Series(energies, spin_factor * dos, name) + + +def _flip_down_component(name): + return "down" in name and "up" not in name and "total" not in name diff --git a/src/py4vasp/_calculation/effective_coulomb.py b/src/py4vasp/_calculation/effective_coulomb.py index 3ab958986..ef60d4652 100644 --- a/src/py4vasp/_calculation/effective_coulomb.py +++ b/src/py4vasp/_calculation/effective_coulomb.py @@ -1,686 +1,686 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from dataclasses import asdict, dataclass -from types import EllipsisType - -import numpy as np -from numpy.typing import ArrayLike - -from py4vasp import exception, interpolate, raw -from py4vasp._calculation.cell import CellHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import EffectiveCoulomb_DB -from py4vasp._third_party import graph, numeric -from py4vasp._util import check, convert, index, select - - -class EffectiveCoulombHandler: - """Handler for the effective_coulomb quantity. Works with exactly one raw.EffectiveCoulomb object.""" - - def __init__(self, raw_coulomb: raw.EffectiveCoulomb): - self._raw_coulomb = raw_coulomb - - @classmethod - def from_data(cls, raw_coulomb: raw.EffectiveCoulomb) -> "EffectiveCoulombHandler": - return cls(raw_coulomb) - - def __str__(self) -> str: - data = asdict(self.to_database()) - return f"""\ -averaged bare interaction -bare Hubbard U = {data["bare_V_uppercase"].real:8.4f} {data["bare_V_uppercase"].imag:8.4f} -bare Hubbard u = {data["bare_v_lowercase"].real:8.4f} {data["bare_v_lowercase"].imag:8.4f} -bare Hubbard J = {data["bare_J_uppercase"].real:8.4f} {data["bare_J_uppercase"].imag:8.4f} - -averaged interaction parameter -screened Hubbard U = {data["screened_U_uppercase"].real:8.4f} {data["screened_U_uppercase"].imag:8.4f} -screened Hubbard u = {data["screened_u_lowercase"].real:8.4f} {data["screened_u_lowercase"].imag:8.4f} -screened Hubbard J = {data["screened_J_uppercase"].real:8.4f} {data["screened_J_uppercase"].imag:8.4f} -""" - - def to_database(self) -> EffectiveCoulomb_DB: - """Serialize effective Coulomb data for database storage.""" - wannier_iiii = self._wannier_indices_iiii() - wannier_ijji = self._wannier_indices_ijji() - wannier_ijij = self._wannier_indices_ijij() - spin_diagonal = slice(None, 2) - omega_0 = origin = 0 - complex_ = slice(None) - if self._has_positions and self._has_frequencies: - access_U = (omega_0, spin_diagonal, wannier_iiii, origin, complex_) - access_u = (omega_0, spin_diagonal, wannier_ijji, origin, complex_) - access_J = (omega_0, spin_diagonal, wannier_ijij, origin, complex_) - access_V = (spin_diagonal, wannier_iiii, origin, complex_) - access_v = (spin_diagonal, wannier_ijji, origin, complex_) - access_Vj = (spin_diagonal, wannier_ijij, origin, complex_) - elif self._has_frequencies: - access_U = (omega_0, spin_diagonal, wannier_iiii, complex_) - access_u = (omega_0, spin_diagonal, wannier_ijji, complex_) - access_J = (omega_0, spin_diagonal, wannier_ijij, complex_) - access_V = (spin_diagonal, wannier_iiii, complex_) - access_v = (spin_diagonal, wannier_ijji, complex_) - access_Vj = (spin_diagonal, wannier_ijij, complex_) - elif self._has_positions: - access_U = access_V = (spin_diagonal, wannier_iiii, origin, complex_) - access_u = access_v = (spin_diagonal, wannier_ijji, origin, complex_) - access_J = access_Vj = (spin_diagonal, wannier_ijij, origin, complex_) - else: - access_U = access_V = (spin_diagonal, wannier_iiii, complex_) - access_u = access_v = (spin_diagonal, wannier_ijji, complex_) - access_J = access_Vj = (spin_diagonal, wannier_ijij, complex_) - U = convert.to_complex(self._raw_coulomb.screened_potential[access_U]) - u = convert.to_complex(self._raw_coulomb.screened_potential[access_u]) - J = convert.to_complex(self._raw_coulomb.screened_potential[access_J]) - V = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[access_V]) - v = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[access_v]) - Vj = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[access_Vj]) - return EffectiveCoulomb_DB( - screened_U_uppercase=complex(np.average(U)), - screened_u_lowercase=complex(np.average(u)), - screened_J_uppercase=complex(np.average(J)), - bare_V_uppercase=complex(np.average(V)), - bare_v_lowercase=complex(np.average(v)), - bare_J_uppercase=complex(np.average(Vj)), - ) - - def _wannier_indices_iiii(self): - n = self._raw_coulomb.number_wannier_states - step = n**3 + n**2 + n + 1 - stop = n**4 - return slice(0, stop, step) - - def _wannier_indices_ijij(self): - n = self._raw_coulomb.number_wannier_states - stop = n**4 - slice_included = slice(0, stop, n**2 + 1) - slice_excluded = slice(0, stop, n**3 + n**2 + n + 1) - indices = np.arange(stop) - return np.setdiff1d(indices[slice_included], indices[slice_excluded]) - - def _wannier_indices_ijji(self): - n = self._raw_coulomb.number_wannier_states - stop = n**4 - indices_included = np.concatenate( - [i * (n**3 + 1) + np.arange(0, n**3, n**2 + n) for i in range(n)] - ) - slice_excluded = slice(0, stop, n**3 + n**2 + n + 1) - indices = np.arange(stop) - return np.setdiff1d(indices_included, indices[slice_excluded]) - - def to_dict(self) -> dict[str, np.ndarray]: - """Convert the effective Coulomb object to a dictionary representation. - - The integrals are evaluated over 4 Wannier functions. For the bare Coulomb - interaction, these integrals can be computed with either a high :tag:`ENCUT` - or low cutoff :tag:`ENCUTGW` that you set in the INCAR file. The screened Coulomb - interaction is evaluated with the dielectric function and will have smaller - values than the bare Coulomb potential. If you set :tag:`TWO_CENTER` = `.TRUE.` - in the INCAR file, the Coulomb interactions are evaluated also at neighboring - cells. - - Returns - ------- - - - A dictionary containing the effective Coulomb interaction data. In particular, - it includes the bare Coulomb interaction with high and low cutoffs, the screened - Coulomb interaction, and optionally the frequencies and positions at which the - interactions are evaluated. - """ - return { - "bare_high_cutoff": self._read_high_cutoff(), - "bare_low_cutoff": self._read_low_cutoff(), - "screened": self._read_screened(), - **self._read_frequencies(), - **self._read_positions(), - } - - @property - def _has_frequencies(self): - return len(self._raw_coulomb.frequencies) > 1 - - @property - def _has_positions(self): - return not check.is_none(self._raw_coulomb.positions) - - @property - def _is_collinear(self): - return len(self._raw_coulomb.bare_potential_low_cutoff) == 3 - - def _read_high_cutoff(self): - V = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[:]) - if self._has_positions: - V = np.moveaxis(V, -1, 0) - V = self._unpack_wannier_indices(V) - if self._has_frequencies: - V = V[..., np.newaxis] - return V - - def _read_low_cutoff(self): - C = convert.to_complex(self._raw_coulomb.bare_potential_low_cutoff[:]) - C = self._unpack_wannier_indices(C) - if self._has_frequencies: - C = C[..., np.newaxis] - return C - - def _read_screened(self): - U = convert.to_complex(self._raw_coulomb.screened_potential[:]) - if self._has_positions: - U = np.moveaxis(U, -1, 0) - U = self._unpack_wannier_indices(U) - if self._has_frequencies: - U = np.moveaxis(U, 1 if self._has_positions else 0, -1) - return U - - def _unpack_wannier_indices(self, data): - num_wannier = self._raw_coulomb.number_wannier_states - new_shape = data.shape[:-1] + 4 * (num_wannier,) - return data.reshape(new_shape) - - def _read_frequencies(self): - if not self._has_frequencies: - return {} - return {"frequencies": convert.to_complex(self._raw_coulomb.frequencies[:])} - - def _read_positions(self): - if not self._has_positions: - return {} - return { - "lattice_vectors": self._cell().lattice_vectors(), - "positions": self._raw_coulomb.positions[:], - } - - def _cell(self): - return CellHandler.from_data(self._raw_coulomb.cell, steps=-1) - - def to_graph( - self, - selection: str = "U J V", - omega: None | EllipsisType | np.ndarray = None, - radius: None | EllipsisType | np.ndarray = None, - radius_max: None | float = None, - config=None, - ) -> graph.Graph: - """Generate a graph representation of the effective Coulomb interaction. - - The method automatically determines the plot type based on which parameters - are provided: - - - If only omega is given: creates a frequency-dependent plot - - If only radius/radius_max is given: creates a radial-dependent plot - - If both omega and radius/radius_max are given: creates a frequency plot for all radii - - Parameters - ---------- - selection - Specifies which data to plot. Default is "U", "V", and "J". You can also - select "u" or "v". For collinear calculations, you select a specific spin - coupling with "up~up" or "up~down". Different choices can be combined, e.g., - "U(up~up)" or "J(up~down)". You may prefix the selection with "bare" or - "screened" to select the bare or screened potential, e.g., "bare(U)". If no - prefix is given, it is deduced from the selection, e.g., "U" is interpreted - as "screened(U)". - - omega - Frequency values for frequency-dependent plots. If not set, or set to - ellipsis (...), the frequency points along the imaginary axis are used. - You can also provide specific frequencies on the real axis, then the data - will be analytically continued and plotted for the selected frequencies. - - radius - Radial distance values for radial-dependent plots. If not set, the plot - will be for r=0. If set to ellipsis (...), the radial points used in VASP - are used. You can also provide specific radii, then the data will be - interpolated to the selected radii. - - radius_max - Maximum radius for radial-dependent plots. If set, all data for radii - greater than this value will be ignored. - - config - Configuration for the analytic continuation of the frequency-dependent data. - Use this if you need to adjust the parameters of the analytic continuation. - - Returns - ------- - - - A graph object containing the visualization of the effective Coulomb - interaction data. - """ - if config is None: - config = interpolate.AAAConfig() - tree = select.Tree.from_selection(selection) - plotter = self._make_plotter(omega, radius, radius_max, config) - potentials = self._get_effective_potentials(tree, plotter) - series = plotter.make_all_series(potentials) - return graph.Graph( - series, xlabel=plotter.xlabel, ylabel="Coulomb potential (eV)" - ) - - def _make_plotter(self, omega, radius, radius_max, config): - radius_set = radius is not None or radius_max is not None - if omega is not None and radius_set: - if radius is not ... and radius is not None: - raise exception.NotImplemented( - "Interpolating radial data for frequency plots is not implemented." - ) - omega_in = self._read_frequencies().get("frequencies") - positions = self._read_positions() - return _OmegaPlotter(omega_in, omega, config, positions, radius_max) - if radius_set: - positions = self._read_positions() - return _RadialPlotter(positions, radius, radius_max) - else: - omega_in = self._read_frequencies().get("frequencies") - return _OmegaPlotter(omega_in, omega, config) - - def _get_effective_potentials(self, tree, plotter): - return [ - self._get_effective_potential(selection, plotter) - for selection in tree.selections() - ] - - def _get_effective_potential(self, selection, plotter): - if self._bare_potential_selected(selection): - return self._get_bare_potential(selection, plotter) - else: - return self._get_screened_potential(selection, plotter) - - def _bare_potential_selected(self, selection): - if select.contains(selection, "bare"): - return True - if select.contains(selection, "screened"): - return False - return select.contains(selection, "V") or select.contains(selection, "v") - - def _get_bare_potential(self, selection, plotter): - selection = self._filter_component_from_selection(selection) - maps = self._create_map("bare") - potential = self._raw_coulomb.bare_potential_high_cutoff - selector = index.Selector(maps, potential, reduction=np.average) - V = convert.to_complex(selector[selection]) - V = plotter.interpolate_bare_if_necessary(V) - return _CoulombPotential("bare", selector.label(selection), V) - - def _get_screened_potential(self, selection, plotter): - selection = self._filter_component_from_selection(selection) - maps = self._create_map("screened") - potential = self._raw_coulomb.screened_potential - selector = index.Selector(maps, potential, reduction=np.average) - U = convert.to_complex(selector[selection]) - U = plotter.interpolate_screened_if_necessary(U) - return _CoulombPotential("screened", selector.label(selection), U) - - def _filter_component_from_selection(self, selection): - return tuple(part for part in selection if part not in {"bare", "screened"}) - - def _create_map(self, component): - spin_map = self._create_spin_map() - component_map = self._create_component_map() - if component == "bare" or not self._has_frequencies: - return {0: spin_map, 1: component_map} - else: - return {1: spin_map, 2: component_map} - - def _create_spin_map(self): - if self._is_collinear: - spin_map = { - convert.text_to_string(label): slice(i, i + 1) - for i, label in enumerate(self._raw_coulomb.spin_labels[:]) - } - spin_map[None] = spin_map["total"] = slice(0, 2) - else: - spin_map = {None: slice(None), "total": slice(None)} - return spin_map - - def _create_component_map(self): - wannier_iiii = self._wannier_indices_iiii() - wannier_ijij = self._wannier_indices_ijij() - wannier_ijji = self._wannier_indices_ijji() - return { - None: wannier_iiii, - "U": wannier_iiii, - "u": wannier_ijji, - "J": wannier_ijij, - "V": wannier_iiii, - "v": wannier_ijji, - } - - def selections(self) -> dict[str, list[str]]: - """Return a dictionary describing what kind of data are available.""" - spin_map = self._create_spin_map() - component_map = self._create_component_map() - return { - "spin": [str(key) for key in spin_map if key is not None], - "screening": ["screened", "bare"], - "potential": [str(key) for key in component_map if key is not None], - } - - -@quantity("effective_coulomb") -class EffectiveCoulomb(graph.Mixin): - """Effective Coulomb interaction U obtained with the constrained random phase approximation (cRPA). - - This class provides post-processing routines to read and visualize first-principles - results from constrained Random Phase Approximation (cRPA) calculations. After you - have performed a cRPA calculation using VASP this class can visualize the effective - Coulomb interaction *U* along the radial or frequency axis. Youy can use this *U* - mean-field theories like DFT+*U* and Dynamical Mean Field Theory (DMFT). - - The cRPA method is essential for strongly correlated materials, where standard Density - Functional Theory (DFT) often incorrectly predicts a metallic ground state or fails to - capture magnetic order. You can activate the cRPA calculation in VASP by setting - :tag:`ALGO` = `CRPAR` in the INCAR file. The method computes the effective Coulomb - interaction *U* in real space by excluding screening processes within a predefined - correlated subspace, typically associated with localized orbitals such as *d* or *f* - states. - - While different flavors of cRPA exist, we recommend using the spectral cRPA (s-cRPA) - method that you activate by setting :tag:`LSCRPA` = `.TRUE.`. in the INCAR file. This - approach overcomes significant limitations of earlier cRPA formulations [1]_, in - particular numerical instabilities for highly occupied correlated shells or unphysical - results like negative *U* values. - - References - ---------- - .. [1] Kaltak, M., *et al.*, Constrained Random Phase Approximation: the spectral - method, Phys. Rev. B 112, 245102 (2025), https://doi.org/10.1103/m3gh-g6r6 - """ - - def __init__(self, source, quantity_name: str = "effective_coulomb"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_coulomb: raw.EffectiveCoulomb) -> "EffectiveCoulomb": - """Create an EffectiveCoulomb dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_coulomb)) - - @property - def path(self): - """Returns the path from which the output is obtained.""" - return self._path - - def _handler_factory(self, raw_data): - return EffectiveCoulombHandler.from_data(raw_data) - - def read(self) -> dict[str, np.ndarray]: - """Convert the effective Coulomb object to a dictionary representation. - - The integrals are evaluated over 4 Wannier functions. For the bare Coulomb - interaction, these integrals can be computed with either a high :tag:`ENCUT` - or low cutoff :tag:`ENCUTGW` that you set in the INCAR file. The screened Coulomb - interaction is evaluated with the dielectric function and will have smaller - values than the bare Coulomb potential. If you set :tag:`TWO_CENTER` = `.TRUE.` - in the INCAR file, the Coulomb interactions are evaluated also at neighboring - cells. - - Returns - ------- - - - A dictionary containing the effective Coulomb interaction data. In particular, - it includes the bare Coulomb interaction with high and low cutoffs, the screened - Coulomb interaction, and optionally the frequencies and positions at which the - interactions are evaluated. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - EffectiveCoulombHandler.to_dict, - ) - - def to_dict(self) -> dict[str, np.ndarray]: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def to_graph( - self, - selection: str = "U J V", - omega: None | EllipsisType | np.ndarray = None, - radius: None | EllipsisType | np.ndarray = None, - radius_max: None | float = None, - config=None, - ) -> graph.Graph: - """Generate a graph representation of the effective Coulomb interaction. - - The method automatically determines the plot type based on which parameters - are provided: - - - If only omega is given: creates a frequency-dependent plot - - If only radius/radius_max is given: creates a radial-dependent plot - - If both omega and radius/radius_max are given: creates a frequency plot for all radii - - Parameters - ---------- - selection - Specifies which data to plot. Default is "U", "V", and "J". You can also - select "u" or "v". For collinear calculations, you select a specific spin - coupling with "up~up" or "up~down". Different choices can be combined, e.g., - "U(up~up)" or "J(up~down)". You may prefix the selection with "bare" or - "screened" to select the bare or screened potential, e.g., "bare(U)". If no - prefix is given, it is deduced from the selection, e.g., "U" is interpreted - as "screened(U)". - - omega - Frequency values for frequency-dependent plots. If not set, or set to - ellipsis (...), the frequency points along the imaginary axis are used. - You can also provide specific frequencies on the real axis, then the data - will be analytically continued and plotted for the selected frequencies. - - radius - Radial distance values for radial-dependent plots. If not set, the plot - will be for r=0. If set to ellipsis (...), the radial points used in VASP - are used. You can also provide specific radii, then the data will be - interpolated to the selected radii. - - radius_max - Maximum radius for radial-dependent plots. If set, all data for radii - greater than this value will be ignored. - - config - Configuration for the analytic continuation of the frequency-dependent data. - Use this if you need to adjust the parameters of the analytic continuation. - - Returns - ------- - - - A graph object containing the visualization of the effective Coulomb - interaction data. - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EffectiveCoulombHandler.to_graph, - omega=omega, - radius=radius, - radius_max=radius_max, - config=config, - ) - - def selections(self) -> dict[str, list[str]]: - """Return a dictionary describing what kind of data are available. - - Returns - ------- - - - Dictionary containing available selection options with their possible values. - Keys include the selection criteria "spin", "screening", and "potential". - """ - result = merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - EffectiveCoulombHandler.selections, - ) - return {"effective_coulomb": ["default"], **result} - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EffectiveCoulombHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - @staticmethod - def ohno_potential(radius: ArrayLike, delta: float) -> np.ndarray: - """Ohno potential for given radius/radii and delta. - - This is used to interpolate the Coulomb potential to other radii. - - Parameters - ---------- - radius - The radial distance(s) at which to evaluate the potential. - delta - The delta parameter for the Ohno potential. - - Returns - ------- - - - The Ohno potential evaluated at the given radius/radii. - """ - delta = np.abs(delta) - return np.sqrt(delta / (radius + delta)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - EffectiveCoulombHandler.from_data, - EffectiveCoulombHandler.to_database, - ) - - -@dataclass -class _CoulombPotential: - component: str - selector_label: str - strength: ArrayLike - - -class _OmegaPlotter: - def __init__(self, omega_in, omega_out, config, positions=None, radius_max=None): - self.omega_in = omega_in - self.interpolate = omega_out is not None and omega_out is not ... - if omega_in is None: - raise exception.DataMismatch("The output does not contain frequency data.") - if positions is not None and not positions: - raise exception.DataMismatch("The output does not contain position data.") - self.omega_out = omega_out if self.interpolate else omega_in - self.xlabel = "ω (eV)" if self.interpolate else "Im(ω) (eV)" - self.config = config - if positions is not None: - self.positions = positions["positions"] - _, self.mask = transform_positions_to_radial(positions, radius_max) - - def interpolate_bare_if_necessary(self, potential): - num_omega = len(self.omega_out) - return np.broadcast_to(potential, (num_omega,) + potential.shape) - - def interpolate_screened_if_necessary(self, potential): - if not self.interpolate: - return potential - return numeric.analytic_continuation( - self.omega_in, - potential.T, - self.omega_out, - config=self.config, - ).T - - def make_all_series(self, potentials): - if hasattr(self, "positions"): - return [ - self._make_one_series(potential, position=position) - for potential in potentials - for position in self.positions[self.mask] - ] - else: - return [self._make_one_series(potential) for potential in potentials] - - def _make_one_series(self, potential, position=None): - if position is None: - position_index = 0 - suffix = "" - else: - lookup = np.all(self.positions == position, axis=1) - position_index = np.squeeze(np.where(lookup)) - suffix = f" @ {position}" - omega = self.omega_out.real if self.interpolate else self.omega_out.imag - label = f"{potential.component} {potential.selector_label}{suffix}" - if potential.strength.ndim == 1: - strength = potential.strength - else: - strength = potential.strength[:, position_index] - return graph.Series(omega, strength.real, label=label) - - -class _RadialPlotter: - xlabel = "Radius (Å)" - - def __init__(self, positions, radius_out, radius_max): - self.radius_in, self.mask = transform_positions_to_radial(positions, radius_max) - self.interpolate = radius_out is not None and radius_out is not ... - self.radius_out = radius_out if self.interpolate else self.radius_in - self.marker = None if self.interpolate else "*" - - def interpolate_bare_if_necessary(self, potential): - return self._ohno_interpolation(potential) - - def interpolate_screened_if_necessary(self, potential): - return self._ohno_interpolation(potential) - - def _ohno_interpolation(self, potential): - potential = potential.real[..., self.mask] - if not self.interpolate: - return potential - if potential.ndim == 2: - # if multiple frequencies are present, take only omega = 0 - potential = potential[0] - return potential[0] * numeric.interpolate_with_function( - EffectiveCoulomb.ohno_potential, - self.radius_in, - potential / potential[0], - self.radius_out, - ) - - def make_all_series(self, potentials): - return [self._make_one_series(potential) for potential in potentials] - - def _make_one_series(self, potential): - label = f"{potential.component} {potential.selector_label}" - if potential.strength.ndim == 1: - strength = potential.strength - else: - # use only omega = 0 - strength = potential.strength[0] - return graph.Series(self.radius_out, strength, label=label, marker=self.marker) - - -def transform_positions_to_radial(positions, radius_max): - if not positions: - raise exception.DataMismatch("The output does not contain position data.") - radius = np.linalg.norm( - positions["lattice_vectors"] @ positions["positions"].T, axis=0 - ) - if radius_max is None: - return radius, slice(None) - mask = radius <= radius_max - return radius[mask], mask +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from dataclasses import asdict, dataclass +from types import EllipsisType + +import numpy as np +from numpy.typing import ArrayLike + +from py4vasp import exception, interpolate, raw +from py4vasp._calculation.cell import CellHandler +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import EffectiveCoulomb_DB +from py4vasp._third_party import graph, numeric +from py4vasp._util import check, convert, index, select + + +class EffectiveCoulombHandler: + """Handler for the effective_coulomb quantity. Works with exactly one raw.EffectiveCoulomb object.""" + + def __init__(self, raw_coulomb: raw.EffectiveCoulomb): + self._raw_coulomb = raw_coulomb + + @classmethod + def from_data(cls, raw_coulomb: raw.EffectiveCoulomb) -> "EffectiveCoulombHandler": + return cls(raw_coulomb) + + def __str__(self) -> str: + data = asdict(self.to_database()) + return f"""\ +averaged bare interaction +bare Hubbard U = {data["bare_V_uppercase"].real:8.4f} {data["bare_V_uppercase"].imag:8.4f} +bare Hubbard u = {data["bare_v_lowercase"].real:8.4f} {data["bare_v_lowercase"].imag:8.4f} +bare Hubbard J = {data["bare_J_uppercase"].real:8.4f} {data["bare_J_uppercase"].imag:8.4f} + +averaged interaction parameter +screened Hubbard U = {data["screened_U_uppercase"].real:8.4f} {data["screened_U_uppercase"].imag:8.4f} +screened Hubbard u = {data["screened_u_lowercase"].real:8.4f} {data["screened_u_lowercase"].imag:8.4f} +screened Hubbard J = {data["screened_J_uppercase"].real:8.4f} {data["screened_J_uppercase"].imag:8.4f} +""" + + def to_database(self) -> EffectiveCoulomb_DB: + """Serialize effective Coulomb data for database storage.""" + wannier_iiii = self._wannier_indices_iiii() + wannier_ijji = self._wannier_indices_ijji() + wannier_ijij = self._wannier_indices_ijij() + spin_diagonal = slice(None, 2) + omega_0 = origin = 0 + complex_ = slice(None) + if self._has_positions and self._has_frequencies: + access_U = (omega_0, spin_diagonal, wannier_iiii, origin, complex_) + access_u = (omega_0, spin_diagonal, wannier_ijji, origin, complex_) + access_J = (omega_0, spin_diagonal, wannier_ijij, origin, complex_) + access_V = (spin_diagonal, wannier_iiii, origin, complex_) + access_v = (spin_diagonal, wannier_ijji, origin, complex_) + access_Vj = (spin_diagonal, wannier_ijij, origin, complex_) + elif self._has_frequencies: + access_U = (omega_0, spin_diagonal, wannier_iiii, complex_) + access_u = (omega_0, spin_diagonal, wannier_ijji, complex_) + access_J = (omega_0, spin_diagonal, wannier_ijij, complex_) + access_V = (spin_diagonal, wannier_iiii, complex_) + access_v = (spin_diagonal, wannier_ijji, complex_) + access_Vj = (spin_diagonal, wannier_ijij, complex_) + elif self._has_positions: + access_U = access_V = (spin_diagonal, wannier_iiii, origin, complex_) + access_u = access_v = (spin_diagonal, wannier_ijji, origin, complex_) + access_J = access_Vj = (spin_diagonal, wannier_ijij, origin, complex_) + else: + access_U = access_V = (spin_diagonal, wannier_iiii, complex_) + access_u = access_v = (spin_diagonal, wannier_ijji, complex_) + access_J = access_Vj = (spin_diagonal, wannier_ijij, complex_) + U = convert.to_complex(self._raw_coulomb.screened_potential[access_U]) + u = convert.to_complex(self._raw_coulomb.screened_potential[access_u]) + J = convert.to_complex(self._raw_coulomb.screened_potential[access_J]) + V = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[access_V]) + v = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[access_v]) + Vj = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[access_Vj]) + return EffectiveCoulomb_DB( + screened_U_uppercase=complex(np.average(U)), + screened_u_lowercase=complex(np.average(u)), + screened_J_uppercase=complex(np.average(J)), + bare_V_uppercase=complex(np.average(V)), + bare_v_lowercase=complex(np.average(v)), + bare_J_uppercase=complex(np.average(Vj)), + ) + + def _wannier_indices_iiii(self): + n = self._raw_coulomb.number_wannier_states + step = n**3 + n**2 + n + 1 + stop = n**4 + return slice(0, stop, step) + + def _wannier_indices_ijij(self): + n = self._raw_coulomb.number_wannier_states + stop = n**4 + slice_included = slice(0, stop, n**2 + 1) + slice_excluded = slice(0, stop, n**3 + n**2 + n + 1) + indices = np.arange(stop) + return np.setdiff1d(indices[slice_included], indices[slice_excluded]) + + def _wannier_indices_ijji(self): + n = self._raw_coulomb.number_wannier_states + stop = n**4 + indices_included = np.concatenate( + [i * (n**3 + 1) + np.arange(0, n**3, n**2 + n) for i in range(n)] + ) + slice_excluded = slice(0, stop, n**3 + n**2 + n + 1) + indices = np.arange(stop) + return np.setdiff1d(indices_included, indices[slice_excluded]) + + def to_dict(self) -> dict[str, np.ndarray]: + """Convert the effective Coulomb object to a dictionary representation. + + The integrals are evaluated over 4 Wannier functions. For the bare Coulomb + interaction, these integrals can be computed with either a high :tag:`ENCUT` + or low cutoff :tag:`ENCUTGW` that you set in the INCAR file. The screened Coulomb + interaction is evaluated with the dielectric function and will have smaller + values than the bare Coulomb potential. If you set :tag:`TWO_CENTER` = `.TRUE.` + in the INCAR file, the Coulomb interactions are evaluated also at neighboring + cells. + + Returns + ------- + - + A dictionary containing the effective Coulomb interaction data. In particular, + it includes the bare Coulomb interaction with high and low cutoffs, the screened + Coulomb interaction, and optionally the frequencies and positions at which the + interactions are evaluated. + """ + return { + "bare_high_cutoff": self._read_high_cutoff(), + "bare_low_cutoff": self._read_low_cutoff(), + "screened": self._read_screened(), + **self._read_frequencies(), + **self._read_positions(), + } + + @property + def _has_frequencies(self): + return len(self._raw_coulomb.frequencies) > 1 + + @property + def _has_positions(self): + return not check.is_none(self._raw_coulomb.positions) + + @property + def _is_collinear(self): + return len(self._raw_coulomb.bare_potential_low_cutoff) == 3 + + def _read_high_cutoff(self): + V = convert.to_complex(self._raw_coulomb.bare_potential_high_cutoff[:]) + if self._has_positions: + V = np.moveaxis(V, -1, 0) + V = self._unpack_wannier_indices(V) + if self._has_frequencies: + V = V[..., np.newaxis] + return V + + def _read_low_cutoff(self): + C = convert.to_complex(self._raw_coulomb.bare_potential_low_cutoff[:]) + C = self._unpack_wannier_indices(C) + if self._has_frequencies: + C = C[..., np.newaxis] + return C + + def _read_screened(self): + U = convert.to_complex(self._raw_coulomb.screened_potential[:]) + if self._has_positions: + U = np.moveaxis(U, -1, 0) + U = self._unpack_wannier_indices(U) + if self._has_frequencies: + U = np.moveaxis(U, 1 if self._has_positions else 0, -1) + return U + + def _unpack_wannier_indices(self, data): + num_wannier = self._raw_coulomb.number_wannier_states + new_shape = data.shape[:-1] + 4 * (num_wannier,) + return data.reshape(new_shape) + + def _read_frequencies(self): + if not self._has_frequencies: + return {} + return {"frequencies": convert.to_complex(self._raw_coulomb.frequencies[:])} + + def _read_positions(self): + if not self._has_positions: + return {} + return { + "lattice_vectors": self._cell().lattice_vectors(), + "positions": self._raw_coulomb.positions[:], + } + + def _cell(self): + return CellHandler.from_data(self._raw_coulomb.cell, steps=-1) + + def to_graph( + self, + selection: str = "U J V", + omega: None | EllipsisType | np.ndarray = None, + radius: None | EllipsisType | np.ndarray = None, + radius_max: None | float = None, + config=None, + ) -> graph.Graph: + """Generate a graph representation of the effective Coulomb interaction. + + The method automatically determines the plot type based on which parameters + are provided: + + - If only omega is given: creates a frequency-dependent plot + - If only radius/radius_max is given: creates a radial-dependent plot + - If both omega and radius/radius_max are given: creates a frequency plot for all radii + + Parameters + ---------- + selection + Specifies which data to plot. Default is "U", "V", and "J". You can also + select "u" or "v". For collinear calculations, you select a specific spin + coupling with "up~up" or "up~down". Different choices can be combined, e.g., + "U(up~up)" or "J(up~down)". You may prefix the selection with "bare" or + "screened" to select the bare or screened potential, e.g., "bare(U)". If no + prefix is given, it is deduced from the selection, e.g., "U" is interpreted + as "screened(U)". + + omega + Frequency values for frequency-dependent plots. If not set, or set to + ellipsis (...), the frequency points along the imaginary axis are used. + You can also provide specific frequencies on the real axis, then the data + will be analytically continued and plotted for the selected frequencies. + + radius + Radial distance values for radial-dependent plots. If not set, the plot + will be for r=0. If set to ellipsis (...), the radial points used in VASP + are used. You can also provide specific radii, then the data will be + interpolated to the selected radii. + + radius_max + Maximum radius for radial-dependent plots. If set, all data for radii + greater than this value will be ignored. + + config + Configuration for the analytic continuation of the frequency-dependent data. + Use this if you need to adjust the parameters of the analytic continuation. + + Returns + ------- + - + A graph object containing the visualization of the effective Coulomb + interaction data. + """ + if config is None: + config = interpolate.AAAConfig() + tree = select.Tree.from_selection(selection) + plotter = self._make_plotter(omega, radius, radius_max, config) + potentials = self._get_effective_potentials(tree, plotter) + series = plotter.make_all_series(potentials) + return graph.Graph( + series, xlabel=plotter.xlabel, ylabel="Coulomb potential (eV)" + ) + + def _make_plotter(self, omega, radius, radius_max, config): + radius_set = radius is not None or radius_max is not None + if omega is not None and radius_set: + if radius is not ... and radius is not None: + raise exception.NotImplemented( + "Interpolating radial data for frequency plots is not implemented." + ) + omega_in = self._read_frequencies().get("frequencies") + positions = self._read_positions() + return _OmegaPlotter(omega_in, omega, config, positions, radius_max) + if radius_set: + positions = self._read_positions() + return _RadialPlotter(positions, radius, radius_max) + else: + omega_in = self._read_frequencies().get("frequencies") + return _OmegaPlotter(omega_in, omega, config) + + def _get_effective_potentials(self, tree, plotter): + return [ + self._get_effective_potential(selection, plotter) + for selection in tree.selections() + ] + + def _get_effective_potential(self, selection, plotter): + if self._bare_potential_selected(selection): + return self._get_bare_potential(selection, plotter) + else: + return self._get_screened_potential(selection, plotter) + + def _bare_potential_selected(self, selection): + if select.contains(selection, "bare"): + return True + if select.contains(selection, "screened"): + return False + return select.contains(selection, "V") or select.contains(selection, "v") + + def _get_bare_potential(self, selection, plotter): + selection = self._filter_component_from_selection(selection) + maps = self._create_map("bare") + potential = self._raw_coulomb.bare_potential_high_cutoff + selector = index.Selector(maps, potential, reduction=np.average) + V = convert.to_complex(selector[selection]) + V = plotter.interpolate_bare_if_necessary(V) + return _CoulombPotential("bare", selector.label(selection), V) + + def _get_screened_potential(self, selection, plotter): + selection = self._filter_component_from_selection(selection) + maps = self._create_map("screened") + potential = self._raw_coulomb.screened_potential + selector = index.Selector(maps, potential, reduction=np.average) + U = convert.to_complex(selector[selection]) + U = plotter.interpolate_screened_if_necessary(U) + return _CoulombPotential("screened", selector.label(selection), U) + + def _filter_component_from_selection(self, selection): + return tuple(part for part in selection if part not in {"bare", "screened"}) + + def _create_map(self, component): + spin_map = self._create_spin_map() + component_map = self._create_component_map() + if component == "bare" or not self._has_frequencies: + return {0: spin_map, 1: component_map} + else: + return {1: spin_map, 2: component_map} + + def _create_spin_map(self): + if self._is_collinear: + spin_map = { + convert.text_to_string(label): slice(i, i + 1) + for i, label in enumerate(self._raw_coulomb.spin_labels[:]) + } + spin_map[None] = spin_map["total"] = slice(0, 2) + else: + spin_map = {None: slice(None), "total": slice(None)} + return spin_map + + def _create_component_map(self): + wannier_iiii = self._wannier_indices_iiii() + wannier_ijij = self._wannier_indices_ijij() + wannier_ijji = self._wannier_indices_ijji() + return { + None: wannier_iiii, + "U": wannier_iiii, + "u": wannier_ijji, + "J": wannier_ijij, + "V": wannier_iiii, + "v": wannier_ijji, + } + + def selections(self) -> dict[str, list[str]]: + """Return a dictionary describing what kind of data are available.""" + spin_map = self._create_spin_map() + component_map = self._create_component_map() + return { + "spin": [str(key) for key in spin_map if key is not None], + "screening": ["screened", "bare"], + "potential": [str(key) for key in component_map if key is not None], + } + + +@quantity("effective_coulomb") +class EffectiveCoulomb(graph.Mixin): + """Effective Coulomb interaction U obtained with the constrained random phase approximation (cRPA). + + This class provides post-processing routines to read and visualize first-principles + results from constrained Random Phase Approximation (cRPA) calculations. After you + have performed a cRPA calculation using VASP this class can visualize the effective + Coulomb interaction *U* along the radial or frequency axis. Youy can use this *U* + mean-field theories like DFT+*U* and Dynamical Mean Field Theory (DMFT). + + The cRPA method is essential for strongly correlated materials, where standard Density + Functional Theory (DFT) often incorrectly predicts a metallic ground state or fails to + capture magnetic order. You can activate the cRPA calculation in VASP by setting + :tag:`ALGO` = `CRPAR` in the INCAR file. The method computes the effective Coulomb + interaction *U* in real space by excluding screening processes within a predefined + correlated subspace, typically associated with localized orbitals such as *d* or *f* + states. + + While different flavors of cRPA exist, we recommend using the spectral cRPA (s-cRPA) + method that you activate by setting :tag:`LSCRPA` = `.TRUE.`. in the INCAR file. This + approach overcomes significant limitations of earlier cRPA formulations [1]_, in + particular numerical instabilities for highly occupied correlated shells or unphysical + results like negative *U* values. + + References + ---------- + .. [1] Kaltak, M., *et al.*, Constrained Random Phase Approximation: the spectral + method, Phys. Rev. B 112, 245102 (2025), https://doi.org/10.1103/m3gh-g6r6 + """ + + def __init__(self, source, quantity_name: str = "effective_coulomb"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_coulomb: raw.EffectiveCoulomb) -> "EffectiveCoulomb": + """Create an EffectiveCoulomb dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_coulomb)) + + @property + def path(self): + """Returns the path from which the output is obtained.""" + return self._path + + def _handler_factory(self, raw_data): + return EffectiveCoulombHandler.from_data(raw_data) + + def read(self) -> dict[str, np.ndarray]: + """Convert the effective Coulomb object to a dictionary representation. + + The integrals are evaluated over 4 Wannier functions. For the bare Coulomb + interaction, these integrals can be computed with either a high :tag:`ENCUT` + or low cutoff :tag:`ENCUTGW` that you set in the INCAR file. The screened Coulomb + interaction is evaluated with the dielectric function and will have smaller + values than the bare Coulomb potential. If you set :tag:`TWO_CENTER` = `.TRUE.` + in the INCAR file, the Coulomb interactions are evaluated also at neighboring + cells. + + Returns + ------- + - + A dictionary containing the effective Coulomb interaction data. In particular, + it includes the bare Coulomb interaction with high and low cutoffs, the screened + Coulomb interaction, and optionally the frequencies and positions at which the + interactions are evaluated. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + EffectiveCoulombHandler.to_dict, + ) + + def to_dict(self) -> dict[str, np.ndarray]: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def to_graph( + self, + selection: str = "U J V", + omega: None | EllipsisType | np.ndarray = None, + radius: None | EllipsisType | np.ndarray = None, + radius_max: None | float = None, + config=None, + ) -> graph.Graph: + """Generate a graph representation of the effective Coulomb interaction. + + The method automatically determines the plot type based on which parameters + are provided: + + - If only omega is given: creates a frequency-dependent plot + - If only radius/radius_max is given: creates a radial-dependent plot + - If both omega and radius/radius_max are given: creates a frequency plot for all radii + + Parameters + ---------- + selection + Specifies which data to plot. Default is "U", "V", and "J". You can also + select "u" or "v". For collinear calculations, you select a specific spin + coupling with "up~up" or "up~down". Different choices can be combined, e.g., + "U(up~up)" or "J(up~down)". You may prefix the selection with "bare" or + "screened" to select the bare or screened potential, e.g., "bare(U)". If no + prefix is given, it is deduced from the selection, e.g., "U" is interpreted + as "screened(U)". + + omega + Frequency values for frequency-dependent plots. If not set, or set to + ellipsis (...), the frequency points along the imaginary axis are used. + You can also provide specific frequencies on the real axis, then the data + will be analytically continued and plotted for the selected frequencies. + + radius + Radial distance values for radial-dependent plots. If not set, the plot + will be for r=0. If set to ellipsis (...), the radial points used in VASP + are used. You can also provide specific radii, then the data will be + interpolated to the selected radii. + + radius_max + Maximum radius for radial-dependent plots. If set, all data for radii + greater than this value will be ignored. + + config + Configuration for the analytic continuation of the frequency-dependent data. + Use this if you need to adjust the parameters of the analytic continuation. + + Returns + ------- + - + A graph object containing the visualization of the effective Coulomb + interaction data. + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EffectiveCoulombHandler.to_graph, + omega=omega, + radius=radius, + radius_max=radius_max, + config=config, + ) + + def selections(self) -> dict[str, list[str]]: + """Return a dictionary describing what kind of data are available. + + Returns + ------- + - + Dictionary containing available selection options with their possible values. + Keys include the selection criteria "spin", "screening", and "potential". + """ + result = merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + EffectiveCoulombHandler.selections, + ) + return {"effective_coulomb": ["default"], **result} + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EffectiveCoulombHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + @staticmethod + def ohno_potential(radius: ArrayLike, delta: float) -> np.ndarray: + """Ohno potential for given radius/radii and delta. + + This is used to interpolate the Coulomb potential to other radii. + + Parameters + ---------- + radius + The radial distance(s) at which to evaluate the potential. + delta + The delta parameter for the Ohno potential. + + Returns + ------- + - + The Ohno potential evaluated at the given radius/radii. + """ + delta = np.abs(delta) + return np.sqrt(delta / (radius + delta)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + EffectiveCoulombHandler.from_data, + EffectiveCoulombHandler.to_database, + ) + + +@dataclass +class _CoulombPotential: + component: str + selector_label: str + strength: ArrayLike + + +class _OmegaPlotter: + def __init__(self, omega_in, omega_out, config, positions=None, radius_max=None): + self.omega_in = omega_in + self.interpolate = omega_out is not None and omega_out is not ... + if omega_in is None: + raise exception.DataMismatch("The output does not contain frequency data.") + if positions is not None and not positions: + raise exception.DataMismatch("The output does not contain position data.") + self.omega_out = omega_out if self.interpolate else omega_in + self.xlabel = "ω (eV)" if self.interpolate else "Im(ω) (eV)" + self.config = config + if positions is not None: + self.positions = positions["positions"] + _, self.mask = transform_positions_to_radial(positions, radius_max) + + def interpolate_bare_if_necessary(self, potential): + num_omega = len(self.omega_out) + return np.broadcast_to(potential, (num_omega,) + potential.shape) + + def interpolate_screened_if_necessary(self, potential): + if not self.interpolate: + return potential + return numeric.analytic_continuation( + self.omega_in, + potential.T, + self.omega_out, + config=self.config, + ).T + + def make_all_series(self, potentials): + if hasattr(self, "positions"): + return [ + self._make_one_series(potential, position=position) + for potential in potentials + for position in self.positions[self.mask] + ] + else: + return [self._make_one_series(potential) for potential in potentials] + + def _make_one_series(self, potential, position=None): + if position is None: + position_index = 0 + suffix = "" + else: + lookup = np.all(self.positions == position, axis=1) + position_index = np.squeeze(np.where(lookup)) + suffix = f" @ {position}" + omega = self.omega_out.real if self.interpolate else self.omega_out.imag + label = f"{potential.component} {potential.selector_label}{suffix}" + if potential.strength.ndim == 1: + strength = potential.strength + else: + strength = potential.strength[:, position_index] + return graph.Series(omega, strength.real, label=label) + + +class _RadialPlotter: + xlabel = "Radius (Å)" + + def __init__(self, positions, radius_out, radius_max): + self.radius_in, self.mask = transform_positions_to_radial(positions, radius_max) + self.interpolate = radius_out is not None and radius_out is not ... + self.radius_out = radius_out if self.interpolate else self.radius_in + self.marker = None if self.interpolate else "*" + + def interpolate_bare_if_necessary(self, potential): + return self._ohno_interpolation(potential) + + def interpolate_screened_if_necessary(self, potential): + return self._ohno_interpolation(potential) + + def _ohno_interpolation(self, potential): + potential = potential.real[..., self.mask] + if not self.interpolate: + return potential + if potential.ndim == 2: + # if multiple frequencies are present, take only omega = 0 + potential = potential[0] + return potential[0] * numeric.interpolate_with_function( + EffectiveCoulomb.ohno_potential, + self.radius_in, + potential / potential[0], + self.radius_out, + ) + + def make_all_series(self, potentials): + return [self._make_one_series(potential) for potential in potentials] + + def _make_one_series(self, potential): + label = f"{potential.component} {potential.selector_label}" + if potential.strength.ndim == 1: + strength = potential.strength + else: + # use only omega = 0 + strength = potential.strength[0] + return graph.Series(self.radius_out, strength, label=label, marker=self.marker) + + +def transform_positions_to_radial(positions, radius_max): + if not positions: + raise exception.DataMismatch("The output does not contain position data.") + radius = np.linalg.norm( + positions["lattice_vectors"] @ positions["positions"].T, axis=0 + ) + if radius_max is None: + return radius, slice(None) + mask = radius <= radius_max + return radius[mask], mask diff --git a/src/py4vasp/_calculation/elastic_modulus.py b/src/py4vasp/_calculation/elastic_modulus.py index 44b1b83a1..8b6d27ed7 100644 --- a/src/py4vasp/_calculation/elastic_modulus.py +++ b/src/py4vasp/_calculation/elastic_modulus.py @@ -1,475 +1,474 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from math import pow -from typing import Optional - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation import structure -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import ElasticModulus_DB -from py4vasp._util import check, error -from py4vasp._util.tensor import symmetry_reduce - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - np.linalg.LinAlgError, - AttributeError, - TypeError, - ValueError, - ZeroDivisionError, -) - - -class ElasticModulusHandler: - """Handler for the elastic modulus quantity. Works with exactly one raw.ElasticModulus object.""" - - def __init__(self, raw_elastic_modulus: raw.ElasticModulus): - self._raw_elastic_modulus = raw_elastic_modulus - - @classmethod - def from_data( - cls, raw_elastic_modulus: raw.ElasticModulus - ) -> "ElasticModulusHandler": - return cls(raw_elastic_modulus) - - def to_dict(self) -> dict: - return { - "clamped_ion": self._raw_elastic_modulus.clamped_ion[:], - "relaxed_ion": self._raw_elastic_modulus.relaxed_ion[:], - } - - def __str__(self) -> str: - return f"""Elastic modulus (kBar) -Direction XX YY ZZ XY YZ ZX --------------------------------------------------------------------------------- -{_elastic_modulus_string(self._raw_elastic_modulus.clamped_ion[:], "clamped-ion")} -{_elastic_modulus_string(self._raw_elastic_modulus.relaxed_ion[:], "relaxed-ion")}""" - - def to_database(self) -> dict: - encountered_errors = {} - error_key = "elastic_modulus:default" - - volume_per_atom = None - ( - bulk_modulus, - shear_modulus, - youngs_modulus, - poisson_ratio, - pugh_ratio, - vickers_hardness, - fracture_toughness, - ) = ([None, None, None] for _ in range(7)) - compact_tensor = [None, None, None] - - if not check.is_none(self._raw_elastic_modulus.structure): - structure_obj = structure.Structure.from_data( - self._raw_elastic_modulus.structure - ) - volume = structure_obj.volume() - num_atoms = structure_obj.number_atoms() - volume_per_atom = volume / num_atoms if num_atoms > 0 else None - - total_tensor, ionic_tensor, electronic_tensor = None, None, None - if not check.is_none(self._raw_elastic_modulus.clamped_ion): - electronic_tensor = self._raw_elastic_modulus.clamped_ion[:] - if not check.is_none(self._raw_elastic_modulus.relaxed_ion): - total_tensor = self._raw_elastic_modulus.relaxed_ion[:] - if electronic_tensor is not None and total_tensor is not None: - ionic_tensor = total_tensor - electronic_tensor - - for idt, tensor in enumerate([total_tensor, ionic_tensor, electronic_tensor]): - voigt_tensor = None - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context=f"to_database.tensor[{idt}]", - ): - if not check.is_none(tensor): - compact_tensor[idt] = symmetry_reduce(symmetry_reduce(tensor).T).T - voigt_tensor = compact_tensor[idt] / 10.0 # converting kbar to GPa - compact_tensor[idt] = ( - list([list(l) for l in compact_tensor[idt]]) - if compact_tensor[idt] is not None - else None - ) - - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context=f"to_database.properties[{idt}]", - ): - ( - bulk_modulus[idt], - shear_modulus[idt], - youngs_modulus[idt], - poisson_ratio[idt], - pugh_ratio[idt], - vickers_hardness[idt], - fracture_toughness[idt], - ) = self._compute_elastic_properties( - voigt_tensor, - volume_per_atom=volume_per_atom, - encountered_errors=encountered_errors, - error_key=error_key, - ) - - return ElasticModulus_DB( - total_3d_tensor=compact_tensor[0], - total_bulk_modulus=bulk_modulus[0], - total_shear_modulus=shear_modulus[0], - total_young_modulus=youngs_modulus[0], - total_poisson_ratio=poisson_ratio[0], - total_pugh_ratio=pugh_ratio[0], - total_vickers_hardness=vickers_hardness[0], - total_fracture_toughness=fracture_toughness[0], - ionic_3d_tensor=compact_tensor[1], - ionic_bulk_modulus=bulk_modulus[1], - ionic_shear_modulus=shear_modulus[1], - ionic_young_modulus=youngs_modulus[1], - ionic_poisson_ratio=poisson_ratio[1], - ionic_pugh_ratio=pugh_ratio[1], - ionic_vickers_hardness=vickers_hardness[1], - ionic_fracture_toughness=fracture_toughness[1], - electronic_3d_tensor=compact_tensor[2], - electronic_bulk_modulus=bulk_modulus[2], - electronic_shear_modulus=shear_modulus[2], - electronic_young_modulus=youngs_modulus[2], - electronic_poisson_ratio=poisson_ratio[2], - electronic_pugh_ratio=pugh_ratio[2], - electronic_vickers_hardness=vickers_hardness[2], - electronic_fracture_toughness=fracture_toughness[2], - ) - - def _compute_elastic_properties( - self, - voigt_tensor: np.ndarray, - volume_per_atom: Optional[float] = None, - *, - encountered_errors: Optional[dict[str, list[str]]] = None, - error_key: Optional[str] = None, - ) -> tuple: - ( - bulk_modulus, - shear_modulus, - youngs_modulus, - poisson_ratio, - pugh_ratio, - vickers_hardness, - fracture_toughness, - ) = (None, None, None, None, None, None, None) - - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="compute_elastic_properties.init", - ): - elastic_tensor = _ElasticTensor.from_array(voigt_tensor) - - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="compute_elastic_properties.vrh", - ): - bulk_modulus, shear_modulus, youngs_modulus, poisson_ratio = ( - elastic_tensor.get_VRH() - ) - - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="compute_elastic_properties.pugh", - ): - if shear_modulus is not None and bulk_modulus is not None: - pugh_ratio = ( - shear_modulus / bulk_modulus - if (bulk_modulus != 0 and shear_modulus != 0) - else 0.0 if shear_modulus == 0 else None - ) - - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="compute_elastic_properties.hardness", - ): - vickers_hardness = elastic_tensor.get_hardness() - - with error.suppress_and_record( - encountered_errors, - error_key, - *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, - context="compute_elastic_properties.fracture", - ): - fracture_toughness = elastic_tensor.get_fracture_toughness( - volume_per_atom - ) - vickers_hardness = elastic_tensor.get_hardness() - - return ( - bulk_modulus, - shear_modulus, - youngs_modulus, - poisson_ratio, - pugh_ratio, - vickers_hardness, - fracture_toughness, - ) - - -@quantity("elastic_modulus") -class ElasticModulus: - """The elastic modulus is the second derivative of the energy with respect to strain. - - The elastic modulus, also known as the modulus of elasticity, is a measure of a - material's stiffness and its ability to deform elastically in response to an - applied force. It quantifies the ratio of stress (force per unit area) to strain - (deformation) in a material within its elastic limit. You can use this class to - extract the elastic modulus of a linear response calculation. There are two - variants of the elastic modulus: (i) in the clamped-ion one, the cell is deformed - but the ions are kept in their positions; (ii) in the relaxed-ion one the - atoms are allowed to relax when the cell is deformed. - """ - - def __init__(self, source, quantity_name: str = "elastic_modulus"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_elastic_modulus: raw.ElasticModulus) -> "ElasticModulus": - """Create an ElasticModulus dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_elastic_modulus)) - - def _handler_factory(self, raw_data): - return ElasticModulusHandler.from_data(raw_data) - - def read(self) -> dict: - """Read the clamped-ion and relaxed-ion elastic modulus into a dictionary. - - Returns - ------- - dict - Contains the level of approximation and its associated elastic modulus. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - ElasticModulusHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ElasticModulusHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - ElasticModulusHandler.from_data, - ElasticModulusHandler.to_database, - ) - - -def _elastic_modulus_string(tensor, label): - compact_tensor = symmetry_reduce(symmetry_reduce(tensor).T).T - line = lambda dir_, vec: dir_ + 6 * " " + " ".join(f"{x:11.4f}" for x in vec) - directions = ("XX", "YY", "ZZ", "XY", "YZ", "ZX") - lines = (line(dir_, vec) for dir_, vec in zip(directions, compact_tensor)) - return f"{label:^79}".rstrip() + "\n" + "\n".join(lines) - - -class _ElasticTensor: - """ - Elastic (stiffness) tensor utilities. - - Supports initialization from: - - 6x6 Voigt matrix - - 21-element upper-triangular row - """ - - # Indices of diagonal elements in flattened upper triangle - _diag_indices = np.array([0, 6, 11, 15, 18, 20]) - - # Common index groups for VRH averages - _bulk_indices = np.array([[0, 0], [1, 1], [2, 2], [0, 1], [1, 2], [0, 2]]) - _shear_indices = np.array( - [ - [0, 0], - [1, 1], - [2, 2], - [0, 1], - [1, 2], - [0, 2], - [3, 3], - [4, 4], - [5, 5], - ] - ) - - def __init__(self, C: np.ndarray): - if C.shape == (6, 6): - self._tensor = C - self._row = None - elif C.shape == (21,): - self._row = C - self._tensor = None - else: - raise ValueError("Input must be a (6,6) or (21,) numpy array.") - - self._compliance = None - - # ---------- Constructors ---------- - - @classmethod - def from_array(cls, C: np.ndarray): - return cls(C) - - # ---------- Tensor / row conversion ---------- - - @staticmethod - def _row_to_tensor(row: np.ndarray) -> np.ndarray: - C = np.zeros((6, 6)) - C[np.diag_indices(6)] = row[_ElasticTensor._diag_indices] - - idx = _ElasticTensor._diag_indices.copy() - for k in range(1, 6): - idx = idx[:-1] + 1 - C += np.diag(row[idx], k=k) + np.diag(row[idx], k=-k) - - return C - - @staticmethod - def _tensor_to_row(C: np.ndarray) -> np.ndarray: - return C[np.triu_indices(6)] - - # ---------- Properties ---------- - - @property - def tensor(self) -> np.ndarray: - if self._tensor is None: - self._tensor = self._row_to_tensor(self._row) - return self._tensor - - @property - def row(self) -> np.ndarray: - if self._row is None: - self._row = self._tensor_to_row(self._tensor) - return self._row - - @property - def compliance_tensor(self) -> np.ndarray: - if self._compliance is None: - self._compliance = np.linalg.inv(self.tensor) - return self._compliance - - # ---------- Internal helpers ---------- - - @staticmethod - def _weighted_sum(tensor, indices, coeffs): - return sum(tensor[i, j] * c for (i, j), c in zip(indices, coeffs)) - - # ---------- Mechanical properties ---------- - - def get_VRH(self): - """ - Voigt-Reuss-Hill averaged elastic moduli. - - Returns - ------- - K, G, E, nu - """ - Kv = (1 / 9) * self._weighted_sum( - self.tensor, - self._bulk_indices, - [1, 1, 1, 2, 2, 2], - ) - - Gv = (1 / 15) * self._weighted_sum( - self.tensor, - self._shear_indices, - [1, 1, 1, -1, -1, -1, 3, 3, 3], - ) - - Kr_inv = self._weighted_sum( - self.compliance_tensor, - self._bulk_indices, - [1, 1, 1, 2, 2, 2], - ) - - Gr_inv = (1 / 15) * self._weighted_sum( - self.compliance_tensor, - self._shear_indices, - [4, 4, 4, -4, -4, -4, 3, 3, 3], - ) - - K = 0.5 * (Kv + 1 / Kr_inv) - G = 0.5 * (Gv + 1 / Gr_inv) - E = 9 * K * G / (3 * K + G) - nu = (3 * K - 2 * G) / (6 * K + 2 * G) - - return K, G, E, nu - - def get_hardness(self): - """ - Empirical Vickers hardness (GPa). - - DOI: 10.1063/1.5113622 - """ - _, _, E, nu = self.get_VRH() - return ( - 0.096 - * E - * (1 - 8.5 * nu + 19.5 * pow(nu, 2)) - / (1 - 7.5 * nu + 12.2 * pow(nu, 2) + 19.6 * pow(nu, 3)) - ) - - def get_fracture_toughness(self, V0): - """ - Empirical fracture toughness (MPa m^1/2). - - DOI: 10.1063/1.5113622 - """ - if check.is_none(V0): - return None - _, _, E, nu = self.get_VRH() - return ( - 1e-2 - * pow(8840, -0.5) - * pow(V0, 1.0 / 6.0) - * pow( - ( - E - * (1.0 - 13.7 * nu + 48.6 * pow(nu, 2.0)) - / (1.0 - 15.2 * nu + 70.2 * pow(nu, 2.0) - 81.5 * pow(nu, 3.0)) - ), - 1.5, - ) - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from math import pow +from typing import Optional + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation import structure +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import ElasticModulus_DB +from py4vasp._util import check, error +from py4vasp._util.tensor import symmetry_reduce + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + np.linalg.LinAlgError, + AttributeError, + TypeError, + ValueError, + ZeroDivisionError, +) + + +class ElasticModulusHandler: + """Handler for the elastic modulus quantity. Works with exactly one raw.ElasticModulus object.""" + + def __init__(self, raw_elastic_modulus: raw.ElasticModulus): + self._raw_elastic_modulus = raw_elastic_modulus + + @classmethod + def from_data( + cls, raw_elastic_modulus: raw.ElasticModulus + ) -> "ElasticModulusHandler": + return cls(raw_elastic_modulus) + + def to_dict(self) -> dict: + return { + "clamped_ion": self._raw_elastic_modulus.clamped_ion[:], + "relaxed_ion": self._raw_elastic_modulus.relaxed_ion[:], + } + + def __str__(self) -> str: + return f"""Elastic modulus (kBar) +Direction XX YY ZZ XY YZ ZX +-------------------------------------------------------------------------------- +{_elastic_modulus_string(self._raw_elastic_modulus.clamped_ion[:], "clamped-ion")} +{_elastic_modulus_string(self._raw_elastic_modulus.relaxed_ion[:], "relaxed-ion")}""" + + def to_database(self) -> dict: + encountered_errors = {} + error_key = "elastic_modulus:default" + + volume_per_atom = None + ( + bulk_modulus, + shear_modulus, + youngs_modulus, + poisson_ratio, + pugh_ratio, + vickers_hardness, + fracture_toughness, + ) = ([None, None, None] for _ in range(7)) + compact_tensor = [None, None, None] + + if not check.is_none(self._raw_elastic_modulus.structure): + structure_obj = structure.Structure.from_data( + self._raw_elastic_modulus.structure + ) + volume = structure_obj.volume() + num_atoms = structure_obj.number_atoms() + volume_per_atom = volume / num_atoms if num_atoms > 0 else None + + total_tensor, ionic_tensor, electronic_tensor = None, None, None + if not check.is_none(self._raw_elastic_modulus.clamped_ion): + electronic_tensor = self._raw_elastic_modulus.clamped_ion[:] + if not check.is_none(self._raw_elastic_modulus.relaxed_ion): + total_tensor = self._raw_elastic_modulus.relaxed_ion[:] + if electronic_tensor is not None and total_tensor is not None: + ionic_tensor = total_tensor - electronic_tensor + + for idt, tensor in enumerate([total_tensor, ionic_tensor, electronic_tensor]): + voigt_tensor = None + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context=f"to_database.tensor[{idt}]", + ): + if not check.is_none(tensor): + compact_tensor[idt] = symmetry_reduce(symmetry_reduce(tensor).T).T + voigt_tensor = compact_tensor[idt] / 10.0 # converting kbar to GPa + compact_tensor[idt] = ( + list([list(l) for l in compact_tensor[idt]]) + if compact_tensor[idt] is not None + else None + ) + + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context=f"to_database.properties[{idt}]", + ): + ( + bulk_modulus[idt], + shear_modulus[idt], + youngs_modulus[idt], + poisson_ratio[idt], + pugh_ratio[idt], + vickers_hardness[idt], + fracture_toughness[idt], + ) = self._compute_elastic_properties( + voigt_tensor, + volume_per_atom=volume_per_atom, + encountered_errors=encountered_errors, + error_key=error_key, + ) + + return ElasticModulus_DB( + total_3d_tensor=compact_tensor[0], + total_bulk_modulus=bulk_modulus[0], + total_shear_modulus=shear_modulus[0], + total_young_modulus=youngs_modulus[0], + total_poisson_ratio=poisson_ratio[0], + total_pugh_ratio=pugh_ratio[0], + total_vickers_hardness=vickers_hardness[0], + total_fracture_toughness=fracture_toughness[0], + ionic_3d_tensor=compact_tensor[1], + ionic_bulk_modulus=bulk_modulus[1], + ionic_shear_modulus=shear_modulus[1], + ionic_young_modulus=youngs_modulus[1], + ionic_poisson_ratio=poisson_ratio[1], + ionic_pugh_ratio=pugh_ratio[1], + ionic_vickers_hardness=vickers_hardness[1], + ionic_fracture_toughness=fracture_toughness[1], + electronic_3d_tensor=compact_tensor[2], + electronic_bulk_modulus=bulk_modulus[2], + electronic_shear_modulus=shear_modulus[2], + electronic_young_modulus=youngs_modulus[2], + electronic_poisson_ratio=poisson_ratio[2], + electronic_pugh_ratio=pugh_ratio[2], + electronic_vickers_hardness=vickers_hardness[2], + electronic_fracture_toughness=fracture_toughness[2], + ) + + def _compute_elastic_properties( + self, + voigt_tensor: np.ndarray, + volume_per_atom: Optional[float] = None, + *, + encountered_errors: Optional[dict[str, list[str]]] = None, + error_key: Optional[str] = None, + ) -> tuple: + ( + bulk_modulus, + shear_modulus, + youngs_modulus, + poisson_ratio, + pugh_ratio, + vickers_hardness, + fracture_toughness, + ) = (None, None, None, None, None, None, None) + + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="compute_elastic_properties.init", + ): + elastic_tensor = _ElasticTensor.from_array(voigt_tensor) + + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="compute_elastic_properties.vrh", + ): + bulk_modulus, shear_modulus, youngs_modulus, poisson_ratio = ( + elastic_tensor.get_VRH() + ) + + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="compute_elastic_properties.pugh", + ): + if shear_modulus is not None and bulk_modulus is not None: + pugh_ratio = ( + shear_modulus / bulk_modulus + if (bulk_modulus != 0 and shear_modulus != 0) + else 0.0 if shear_modulus == 0 else None + ) + + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="compute_elastic_properties.hardness", + ): + vickers_hardness = elastic_tensor.get_hardness() + + with error.suppress_and_record( + encountered_errors, + error_key, + *_TO_DATABASE_SUPPRESSED_EXCEPTIONS, + context="compute_elastic_properties.fracture", + ): + fracture_toughness = elastic_tensor.get_fracture_toughness( + volume_per_atom + ) + + return ( + bulk_modulus, + shear_modulus, + youngs_modulus, + poisson_ratio, + pugh_ratio, + vickers_hardness, + fracture_toughness, + ) + + +@quantity("elastic_modulus") +class ElasticModulus: + """The elastic modulus is the second derivative of the energy with respect to strain. + + The elastic modulus, also known as the modulus of elasticity, is a measure of a + material's stiffness and its ability to deform elastically in response to an + applied force. It quantifies the ratio of stress (force per unit area) to strain + (deformation) in a material within its elastic limit. You can use this class to + extract the elastic modulus of a linear response calculation. There are two + variants of the elastic modulus: (i) in the clamped-ion one, the cell is deformed + but the ions are kept in their positions; (ii) in the relaxed-ion one the + atoms are allowed to relax when the cell is deformed. + """ + + def __init__(self, source, quantity_name: str = "elastic_modulus"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_elastic_modulus: raw.ElasticModulus) -> "ElasticModulus": + """Create an ElasticModulus dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_elastic_modulus)) + + def _handler_factory(self, raw_data): + return ElasticModulusHandler.from_data(raw_data) + + def read(self) -> dict: + """Read the clamped-ion and relaxed-ion elastic modulus into a dictionary. + + Returns + ------- + dict + Contains the level of approximation and its associated elastic modulus. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + ElasticModulusHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ElasticModulusHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + ElasticModulusHandler.from_data, + ElasticModulusHandler.to_database, + ) + + +def _elastic_modulus_string(tensor, label): + compact_tensor = symmetry_reduce(symmetry_reduce(tensor).T).T + line = lambda dir_, vec: dir_ + 6 * " " + " ".join(f"{x:11.4f}" for x in vec) + directions = ("XX", "YY", "ZZ", "XY", "YZ", "ZX") + lines = (line(dir_, vec) for dir_, vec in zip(directions, compact_tensor)) + return f"{label:^79}".rstrip() + "\n" + "\n".join(lines) + + +class _ElasticTensor: + """ + Elastic (stiffness) tensor utilities. + + Supports initialization from: + - 6x6 Voigt matrix + - 21-element upper-triangular row + """ + + # Indices of diagonal elements in flattened upper triangle + _diag_indices = np.array([0, 6, 11, 15, 18, 20]) + + # Common index groups for VRH averages + _bulk_indices = np.array([[0, 0], [1, 1], [2, 2], [0, 1], [1, 2], [0, 2]]) + _shear_indices = np.array( + [ + [0, 0], + [1, 1], + [2, 2], + [0, 1], + [1, 2], + [0, 2], + [3, 3], + [4, 4], + [5, 5], + ] + ) + + def __init__(self, C: np.ndarray): + if C.shape == (6, 6): + self._tensor = C + self._row = None + elif C.shape == (21,): + self._row = C + self._tensor = None + else: + raise ValueError("Input must be a (6,6) or (21,) numpy array.") + + self._compliance = None + + # ---------- Constructors ---------- + + @classmethod + def from_array(cls, C: np.ndarray): + return cls(C) + + # ---------- Tensor / row conversion ---------- + + @staticmethod + def _row_to_tensor(row: np.ndarray) -> np.ndarray: + C = np.zeros((6, 6)) + C[np.diag_indices(6)] = row[_ElasticTensor._diag_indices] + + idx = _ElasticTensor._diag_indices.copy() + for k in range(1, 6): + idx = idx[:-1] + 1 + C += np.diag(row[idx], k=k) + np.diag(row[idx], k=-k) + + return C + + @staticmethod + def _tensor_to_row(C: np.ndarray) -> np.ndarray: + return C[np.triu_indices(6)] + + # ---------- Properties ---------- + + @property + def tensor(self) -> np.ndarray: + if self._tensor is None: + self._tensor = self._row_to_tensor(self._row) + return self._tensor + + @property + def row(self) -> np.ndarray: + if self._row is None: + self._row = self._tensor_to_row(self._tensor) + return self._row + + @property + def compliance_tensor(self) -> np.ndarray: + if self._compliance is None: + self._compliance = np.linalg.inv(self.tensor) + return self._compliance + + # ---------- Internal helpers ---------- + + @staticmethod + def _weighted_sum(tensor, indices, coeffs): + return sum(tensor[i, j] * c for (i, j), c in zip(indices, coeffs)) + + # ---------- Mechanical properties ---------- + + def get_VRH(self): + """ + Voigt-Reuss-Hill averaged elastic moduli. + + Returns + ------- + K, G, E, nu + """ + Kv = (1 / 9) * self._weighted_sum( + self.tensor, + self._bulk_indices, + [1, 1, 1, 2, 2, 2], + ) + + Gv = (1 / 15) * self._weighted_sum( + self.tensor, + self._shear_indices, + [1, 1, 1, -1, -1, -1, 3, 3, 3], + ) + + Kr_inv = self._weighted_sum( + self.compliance_tensor, + self._bulk_indices, + [1, 1, 1, 2, 2, 2], + ) + + Gr_inv = (1 / 15) * self._weighted_sum( + self.compliance_tensor, + self._shear_indices, + [4, 4, 4, -4, -4, -4, 3, 3, 3], + ) + + K = 0.5 * (Kv + 1 / Kr_inv) + G = 0.5 * (Gv + 1 / Gr_inv) + E = 9 * K * G / (3 * K + G) + nu = (3 * K - 2 * G) / (6 * K + 2 * G) + + return K, G, E, nu + + def get_hardness(self): + """ + Empirical Vickers hardness (GPa). + + DOI: 10.1063/1.5113622 + """ + _, _, E, nu = self.get_VRH() + return ( + 0.096 + * E + * (1 - 8.5 * nu + 19.5 * pow(nu, 2)) + / (1 - 7.5 * nu + 12.2 * pow(nu, 2) + 19.6 * pow(nu, 3)) + ) + + def get_fracture_toughness(self, V0): + """ + Empirical fracture toughness (MPa m^1/2). + + DOI: 10.1063/1.5113622 + """ + if check.is_none(V0): + return None + _, _, E, nu = self.get_VRH() + return ( + 1e-2 + * pow(8840, -0.5) + * pow(V0, 1.0 / 6.0) + * pow( + ( + E + * (1.0 - 13.7 * nu + 48.6 * pow(nu, 2.0)) + / (1.0 - 15.2 * nu + 70.2 * pow(nu, 2.0) - 81.5 * pow(nu, 3.0)) + ), + 1.5, + ) + ) diff --git a/src/py4vasp/_calculation/electronic_minimization.py b/src/py4vasp/_calculation/electronic_minimization.py index 487a8f54c..ba79ff8b9 100644 --- a/src/py4vasp/_calculation/electronic_minimization.py +++ b/src/py4vasp/_calculation/electronic_minimization.py @@ -1,305 +1,305 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) - -import copy -from contextlib import suppress - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import ElectronicMinimization_DB -from py4vasp._third_party import graph -from py4vasp._util import check - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - AttributeError, - TypeError, - ValueError, - IndexError, -) - - -class ElectronicMinimizationHandler: - """Handler for electronic minimization data — all data access and transformation.""" - - def __init__(self, raw_elmin: raw.ElectronicMinimization, steps=None): - self._raw_data = raw_elmin - self._steps = steps - - @classmethod - def from_data( - cls, raw_elmin: raw.ElectronicMinimization, steps=None - ) -> "ElectronicMinimizationHandler": - return cls(raw_elmin, steps) - - def __str__(self) -> str: - format_rep = "{0:g}\t{1:0.12E}\t{2:0.6E}\t{3:0.6E}\t{4:g}\t{5:0.3E}\t{6:0.3E}\n" - label_rep = "{}\t\t{}\t\t{}\t\t{}\t\t{}\t{}\t\t{}\n" - string = "" - labels = [label.decode("utf-8") for label in getattr(self._raw_data, "label")] - data = self.to_dict() - electronic_iterations = data["N"] - if not self._more_than_one_ionic_step(electronic_iterations): - electronic_iterations = [electronic_iterations] - ionic_steps = len(electronic_iterations) - for ionic_step in range(ionic_steps): - string += label_rep.format(*labels) - electronic_steps = len(electronic_iterations[ionic_step]) - for electronic_step in range(electronic_steps): - _data = [] - for label in self._raw_data.label: - _values_electronic = data[label.decode("utf-8")] - if not self._more_than_one_ionic_step(_values_electronic): - _values_electronic = [_values_electronic] - _value = _values_electronic[ionic_step][electronic_step] - _data.append(_value) - _data = [float(_value) for _value in _data] - string += format_rep.format(*_data) - return string - - def to_dict(self, selection=None) -> dict: - """Extract convergence data and return as a dict.""" - return_data = {} - if selection is None: - keys_to_include = self._from_bytes_to_utf(self._raw_data.label) - else: - labels_as_str = self._from_bytes_to_utf(self._raw_data.label) - if selection not in labels_as_str: - message = """\ -Please choose a selection including at least one of the following keywords: -N, E, dE, deps, ncg, rms, rms(c)""" - raise exception.RefinementError(message) - keys_to_include = [selection] - for key in keys_to_include: - return_data[key] = self._read(key) - return return_data - - def to_graph(self, selection="E") -> graph.Graph: - """Graph the change in parameter with iteration number.""" - data = self.to_dict() - series = graph.Series(data["N"], data[selection], selection) - from py4vasp._util import select as sel_util - - ylabel = " ".join(s.capitalize() for s in selection.split("_")) - return graph.Graph( - series=[series], - xlabel="Iteration number", - ylabel=ylabel, - ) - - def is_converged(self) -> np.ndarray: - is_elmin_converged = self._raw_data.is_elmin_converged[self._steps_or_last] - converged = is_elmin_converged == 0 - if isinstance(converged, bool): - converged = np.array([converged]) - return converged.flatten() - - def to_database(self) -> dict: - """Serialize electronic minimization data for database storage.""" - num_max_electronic_steps_per_ionic = None - num_min_electronic_steps_per_ionic = None - num_electronic_steps = None - elmin_is_converged_all = None - elmin_is_converged_final = None - - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - if not check.is_none(self._raw_data.is_elmin_converged): - elmin_is_converged_all = bool( - np.all(np.array(self._raw_data.is_elmin_converged[:]) == 0.0) - ) - elmin_is_converged_final = bool( - self._raw_data.is_elmin_converged[-1] == 0.0 - ) - - with suppress(exception.NoData, *_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - ( - num_max_electronic_steps_per_ionic, - num_min_electronic_steps_per_ionic, - num_electronic_steps, - ) = self._get_electronic_steps_info() - - return ElectronicMinimization_DB( - num_electronic_steps=num_electronic_steps, - elmin_is_converged_all=elmin_is_converged_all, - elmin_is_converged_final=elmin_is_converged_final, - num_max_electronic_steps_per_ionic_step=num_max_electronic_steps_per_ionic, - num_min_electronic_steps_per_ionic_step=num_min_electronic_steps_per_ionic, - ) - - @property - def _steps_or_last(self): - return -1 if self._steps is None else self._steps - - def _more_than_one_ionic_step(self, data): - return any(isinstance(_data, list) for _data in data) - - def _read(self, key): - data = getattr(self._raw_data, "convergence_data") - iteration_number = data[:, 0] - split_index = np.where(iteration_number == 1)[0] - data = np.vsplit(data, split_index)[1:][self._steps_or_last] - if isinstance(self._steps, slice): - data = [raw.VaspData(_data) for _data in data] - else: - data = [raw.VaspData(data)] - labels = [label.decode("utf-8") for label in self._raw_data.label] - data_index = labels.index(key) - return_data = [list(_data[:, data_index]) for _data in data] - is_none = [_data.is_none() for _data in data] - if len(return_data) == 1: - return_data = return_data[0] - return return_data if not np.all(is_none) else {} - - def _get_electronic_steps_info(self) -> tuple: - if check.is_none(self._raw_data.convergence_data): - return None, None, None - - data = getattr(self._raw_data, "convergence_data") - iteration_number = data[:, 0] - split_index = np.where(iteration_number == 1)[0] - data = [raw.VaspData(_data) for _data in np.vsplit(data, split_index)[1:][:]] - - labels = [label.decode("utf-8") for label in self._raw_data.label] - data_index = labels.index("N") - N_data = [list(_data[:, data_index]) for _data in data] - num_electronic_steps_per_ionic = [len(_data) for _data in N_data] - is_none = [_data.is_none() for _data in data] - if np.all(is_none): - return None, None, None - - if len(num_electronic_steps_per_ionic) == 0: - return None, None, None - - num_max_electronic_steps_per_ionic = max(num_electronic_steps_per_ionic) - num_min_electronic_steps_per_ionic = min(num_electronic_steps_per_ionic) - num_electronic_steps = sum(num_electronic_steps_per_ionic) - - return ( - num_max_electronic_steps_per_ionic, - num_min_electronic_steps_per_ionic, - num_electronic_steps, - ) - - def _from_bytes_to_utf(self, quantity): - return [_quantity.decode("utf-8") for _quantity in quantity] - - -@quantity("electronic_minimization") -class ElectronicMinimization(graph.Mixin): - """Access the convergence data for each electronic step. - - The OSZICAR file written out by VASP stores information related to convergence. - Please check the `vasp-wiki `__ for more - details about the exact outputs generated for each combination of INCAR tags.""" - - def __init__( - self, source, quantity_name: str = "electronic_minimization", steps=None - ): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_elmin: raw.ElectronicMinimization): - """Create an ElectronicMinimization dispatcher from raw data.""" - return cls(source=DataSource(raw_elmin)) - - def __getitem__(self, steps) -> "ElectronicMinimization": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw_data): - return ElectronicMinimizationHandler.from_data(raw_data, steps=self._steps) - - def read(self, selection=None) -> dict: - """Extract convergence data from the HDF5 file and make it available in a dict - - Parameters - ---------- - selection: str - Choose from either iteration_number, free_energy, free_energy_change, - bandstructure_energy_change, number_hamiltonian_evaluations, norm_residual, - difference_charge_density to get specific columns of the OSZICAR file. In - case no selection is provided, supply all columns. - - Returns - ------- - dict - Contains a dict from the HDF5 related to OSZICAR convergence data - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ElectronicMinimizationHandler.to_dict, - ) - - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - def to_graph(self, selection="E") -> graph.Graph: - """Graph the change in parameter with iteration number. - - Parameters - ---------- - selection: str - Choose strings consistent with the OSZICAR format - - Returns - ------- - Graph - The Graph with the quantity plotted on y-axis and the iteration number of - the x-axis. - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ElectronicMinimizationHandler.to_graph, - ) - - def is_converged(self) -> np.ndarray: - """Return whether the electronic minimization converged.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - ElectronicMinimizationHandler.is_converged, - ) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ElectronicMinimizationHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - ElectronicMinimizationHandler.from_data, - ElectronicMinimizationHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import copy +from contextlib import suppress + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import ElectronicMinimization_DB +from py4vasp._third_party import graph +from py4vasp._util import check + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + AttributeError, + TypeError, + ValueError, + IndexError, +) + + +class ElectronicMinimizationHandler: + """Handler for electronic minimization data — all data access and transformation.""" + + def __init__(self, raw_elmin: raw.ElectronicMinimization, steps=None): + self._raw_data = raw_elmin + self._steps = steps + + @classmethod + def from_data( + cls, raw_elmin: raw.ElectronicMinimization, steps=None + ) -> "ElectronicMinimizationHandler": + return cls(raw_elmin, steps) + + def __str__(self) -> str: + format_rep = "{0:g}\t{1:0.12E}\t{2:0.6E}\t{3:0.6E}\t{4:g}\t{5:0.3E}\t{6:0.3E}\n" + label_rep = "{}\t\t{}\t\t{}\t\t{}\t\t{}\t{}\t\t{}\n" + string = "" + labels = [label.decode("utf-8") for label in getattr(self._raw_data, "label")] + data = self.to_dict() + electronic_iterations = data["N"] + if not self._more_than_one_ionic_step(electronic_iterations): + electronic_iterations = [electronic_iterations] + ionic_steps = len(electronic_iterations) + for ionic_step in range(ionic_steps): + string += label_rep.format(*labels) + electronic_steps = len(electronic_iterations[ionic_step]) + for electronic_step in range(electronic_steps): + _data = [] + for label in self._raw_data.label: + _values_electronic = data[label.decode("utf-8")] + if not self._more_than_one_ionic_step(_values_electronic): + _values_electronic = [_values_electronic] + _value = _values_electronic[ionic_step][electronic_step] + _data.append(_value) + _data = [float(_value) for _value in _data] + string += format_rep.format(*_data) + return string + + def to_dict(self, selection=None) -> dict: + """Extract convergence data and return as a dict.""" + return_data = {} + if selection is None: + keys_to_include = self._from_bytes_to_utf(self._raw_data.label) + else: + labels_as_str = self._from_bytes_to_utf(self._raw_data.label) + if selection not in labels_as_str: + message = """\ +Please choose a selection including at least one of the following keywords: +N, E, dE, deps, ncg, rms, rms(c)""" + raise exception.RefinementError(message) + keys_to_include = [selection] + for key in keys_to_include: + return_data[key] = self._read(key) + return return_data + + def to_graph(self, selection="E") -> graph.Graph: + """Graph the change in parameter with iteration number.""" + data = self.to_dict() + series = graph.Series(data["N"], data[selection], selection) + from py4vasp._util import select as sel_util + + ylabel = " ".join(s.capitalize() for s in selection.split("_")) + return graph.Graph( + series=[series], + xlabel="Iteration number", + ylabel=ylabel, + ) + + def is_converged(self) -> np.ndarray: + is_elmin_converged = self._raw_data.is_elmin_converged[self._steps_or_last] + converged = is_elmin_converged == 0 + if isinstance(converged, bool): + converged = np.array([converged]) + return converged.flatten() + + def to_database(self) -> dict: + """Serialize electronic minimization data for database storage.""" + num_max_electronic_steps_per_ionic = None + num_min_electronic_steps_per_ionic = None + num_electronic_steps = None + elmin_is_converged_all = None + elmin_is_converged_final = None + + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + if not check.is_none(self._raw_data.is_elmin_converged): + elmin_is_converged_all = bool( + np.all(np.array(self._raw_data.is_elmin_converged[:]) == 0.0) + ) + elmin_is_converged_final = bool( + self._raw_data.is_elmin_converged[-1] == 0.0 + ) + + with suppress(exception.NoData, *_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + ( + num_max_electronic_steps_per_ionic, + num_min_electronic_steps_per_ionic, + num_electronic_steps, + ) = self._get_electronic_steps_info() + + return ElectronicMinimization_DB( + num_electronic_steps=num_electronic_steps, + elmin_is_converged_all=elmin_is_converged_all, + elmin_is_converged_final=elmin_is_converged_final, + num_max_electronic_steps_per_ionic_step=num_max_electronic_steps_per_ionic, + num_min_electronic_steps_per_ionic_step=num_min_electronic_steps_per_ionic, + ) + + @property + def _steps_or_last(self): + return -1 if self._steps is None else self._steps + + def _more_than_one_ionic_step(self, data): + return any(isinstance(_data, list) for _data in data) + + def _read(self, key): + data = getattr(self._raw_data, "convergence_data") + iteration_number = data[:, 0] + split_index = np.where(iteration_number == 1)[0] + data = np.vsplit(data, split_index)[1:][self._steps_or_last] + if isinstance(self._steps, slice): + data = [raw.VaspData(_data) for _data in data] + else: + data = [raw.VaspData(data)] + labels = [label.decode("utf-8") for label in self._raw_data.label] + data_index = labels.index(key) + return_data = [list(_data[:, data_index]) for _data in data] + is_none = [_data.is_none() for _data in data] + if len(return_data) == 1: + return_data = return_data[0] + return return_data if not np.all(is_none) else {} + + def _get_electronic_steps_info(self) -> tuple: + if check.is_none(self._raw_data.convergence_data): + return None, None, None + + data = getattr(self._raw_data, "convergence_data") + iteration_number = data[:, 0] + split_index = np.where(iteration_number == 1)[0] + data = [raw.VaspData(_data) for _data in np.vsplit(data, split_index)[1:][:]] + + labels = [label.decode("utf-8") for label in self._raw_data.label] + data_index = labels.index("N") + N_data = [list(_data[:, data_index]) for _data in data] + num_electronic_steps_per_ionic = [len(_data) for _data in N_data] + is_none = [_data.is_none() for _data in data] + if np.all(is_none): + return None, None, None + + if len(num_electronic_steps_per_ionic) == 0: + return None, None, None + + num_max_electronic_steps_per_ionic = max(num_electronic_steps_per_ionic) + num_min_electronic_steps_per_ionic = min(num_electronic_steps_per_ionic) + num_electronic_steps = sum(num_electronic_steps_per_ionic) + + return ( + num_max_electronic_steps_per_ionic, + num_min_electronic_steps_per_ionic, + num_electronic_steps, + ) + + def _from_bytes_to_utf(self, quantity): + return [_quantity.decode("utf-8") for _quantity in quantity] + + +@quantity("electronic_minimization") +class ElectronicMinimization(graph.Mixin): + """Access the convergence data for each electronic step. + + The OSZICAR file written out by VASP stores information related to convergence. + Please check the `vasp-wiki `__ for more + details about the exact outputs generated for each combination of INCAR tags.""" + + def __init__( + self, source, quantity_name: str = "electronic_minimization", steps=None + ): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_elmin: raw.ElectronicMinimization): + """Create an ElectronicMinimization dispatcher from raw data.""" + return cls(source=DataSource(raw_elmin)) + + def __getitem__(self, steps) -> "ElectronicMinimization": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw_data): + return ElectronicMinimizationHandler.from_data(raw_data, steps=self._steps) + + def read(self, selection=None) -> dict: + """Extract convergence data from the HDF5 file and make it available in a dict + + Parameters + ---------- + selection: str + Choose from either iteration_number, free_energy, free_energy_change, + bandstructure_energy_change, number_hamiltonian_evaluations, norm_residual, + difference_charge_density to get specific columns of the OSZICAR file. In + case no selection is provided, supply all columns. + + Returns + ------- + dict + Contains a dict from the HDF5 related to OSZICAR convergence data + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ElectronicMinimizationHandler.to_dict, + ) + + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + def to_graph(self, selection="E") -> graph.Graph: + """Graph the change in parameter with iteration number. + + Parameters + ---------- + selection: str + Choose strings consistent with the OSZICAR format + + Returns + ------- + Graph + The Graph with the quantity plotted on y-axis and the iteration number of + the x-axis. + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ElectronicMinimizationHandler.to_graph, + ) + + def is_converged(self) -> np.ndarray: + """Return whether the electronic minimization converged.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + ElectronicMinimizationHandler.is_converged, + ) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ElectronicMinimizationHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + ElectronicMinimizationHandler.from_data, + ElectronicMinimizationHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/energy.py b/src/py4vasp/_calculation/energy.py index 3911ce41b..5d8d6711e 100644 --- a/src/py4vasp/_calculation/energy.py +++ b/src/py4vasp/_calculation/energy.py @@ -1,392 +1,392 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy - -import numpy as np - -from py4vasp import raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import Energy_DB -from py4vasp._third_party import graph -from py4vasp._util import convert, documentation, index, select - - -def _selection_string(default): - return f"""\ -selection : str or None - String specifying the labels of the energy to be read. If no energy is selected - this will default to selecting {default}. Separate distinct labels by commas or - whitespace. You can add or subtract different contributions e.g. `TOTEN + EKIN`. - For a complete list of all possible selections, please use - - >>> calculation.energy.selections() -""" - - -_SELECTIONS = { - "ion-electron TOTEN": ["ion_electron", "TOTEN"], - "kinetic energy EKIN": ["kinetic_energy", "EKIN"], - "kin. lattice EKIN_LAT": ["kinetic_lattice", "EKIN_LAT"], - "temperature TEIN": ["temperature", "TEIN"], - "nose potential ES": ["nose_potential", "ES"], - "nose kinetic EPS": ["nose_kinetic", "EPS"], - "total energy ETOTAL": ["total_energy", "ETOTAL"], - "free energy TOTEN": ["free_energy", "TOTEN"], - "energy without entropy": ["without_entropy", "ENOENT"], - "energy(sigma->0)": ["sigma_0", "ESIG0"], - "step STEP": ["step", "STEP"], - "One el. energy E1": ["one_electron", "E1"], - "Hartree energy -DENC": ["Hartree", "hartree", "DENC"], - "exchange EXHF": ["exchange", "EXHF"], - "free energy TOTEN": ["free_energy", "TOTEN"], - "free energy cap TOTENCAP": ["cap", "TOTENCAP"], - "weight WEIGHT": ["weight", "WEIGHT"], -} - -_DB_KEYS = { - "ion-electron TOTEN": "ion_electron", - "kinetic energy EKIN": "kinetic_energy", - "kin. lattice EKIN_LAT": "kinetic_energy_lattice", - "temperature TEIN": "temperature", - "nose potential ES": "nose_potential", - "nose kinetic EPS": "nose_kinetic", - "total energy ETOTAL": "total_energy", - "free energy TOTEN": "free_energy", - "energy without entropy": "energy_without_entropy", - "energy(sigma->0)": "energy_sigma_0", - "step STEP": "step", - "One el. energy E1": "one_electron_energy", - "Hartree energy -DENC": "hartree_energy", - "exchange EXHF": "exchange_energy", - "free energy TOTEN": "free_energy", - "free energy cap TOTENCAP": "free_energy_cap", - "weight WEIGHT": "weight", -} - - -@documentation.format(examples=slice_.examples("energy")) -class EnergyHandler: - """Handler for energy data — performs all data access and transformation logic.""" - - def __init__(self, raw_energy: raw.Energy, steps=None): - self._raw_energy = raw_energy - self._steps = steps - - @classmethod - def from_data(cls, raw_energy: raw.Energy, steps=None) -> "EnergyHandler": - return cls(raw_energy, steps) - - def __str__(self) -> str: - text = f"Energies at {self._step_string()}:" - values = self._raw_energy.values[self._last_step_in_slice] - for label, value in zip(self._raw_energy.labels, values): - label_str = f"{convert.text_to_string(label):23.23}" - text += f"\n {label_str}={value:17.6f}" - return text - - def to_dict(self, selection=None) -> dict: - if selection is None: - return self._default_dict() - tree = select.Tree.from_selection(selection) - return dict(self._read_data(tree, self._steps_or_last)) - - def to_graph(self, selection="TOTEN") -> graph.Graph: - tree = select.Tree.from_selection(selection) - yaxes = _YAxes(tree) - return graph.Graph( - series=self._make_series(yaxes, tree), - xlabel="Step", - ylabel=yaxes.ylabel, - y2label=yaxes.y2label, - ) - - def to_numpy(self, selection="TOTEN") -> np.ndarray: - tree = select.Tree.from_selection(selection) - return np.squeeze( - [values for _, values in self._read_data(tree, self._steps_or_last)] - ) - - def selections(self) -> dict: - components = list(self._init_selection_dict().keys()) - return {"energy": [], "component": components} - - def to_database(self) -> dict: - default_dict = self._default_dict_all() - energy_dict = {} - for original_label, db_key in _DB_KEYS.items(): - v = default_dict.get(original_label, None) - if v is not None: - v = np.array(v) - energy_dict[f"{db_key}_initial"] = None if v is None else float(v[0]) - if (db_key != "step") and (v is not None): - energy_dict[f"{db_key}_min"] = float(np.min(v)) - energy_dict[f"{db_key}_step_min"] = int(np.argmin(v)) - energy_dict[f"{db_key}_final"] = None if v is None else float(v[-1]) - extra_dict = {} - for k, v in default_dict.items(): - if k not in _DB_KEYS: - vs = np.array(v) if v is not None else None - key = convert.text_to_string(k).strip().lower().replace(" ", "_") - extra_dict[f"{key}_initial"] = None if vs is None else float(vs[0]) - extra_dict[f"{key}_min"] = None if vs is None else float(np.min(vs)) - extra_dict[f"{key}_step_min"] = ( - None if vs is None else int(np.argmin(vs)) - ) - extra_dict[f"{key}_final"] = None if vs is None else float(vs[-1]) - return Energy_DB(**energy_dict, other_energy_data=extra_dict) - - @property - def _steps_or_last(self): - if self._steps is None: - return -1 - return self._steps - - @property - def _last_step_in_slice(self): - if self._steps is None or self._steps == -1: - return -1 - if isinstance(self._steps, slice): - return (self._steps.stop or 0) - 1 - return self._steps - - @property - def _to_slice(self): - if self._steps is None or self._steps == -1: - return slice(-1, None) - if isinstance(self._steps, slice): - return self._steps - return slice(self._steps, self._steps + 1) - - def _step_string(self): - if isinstance(self._steps, slice): - n = len(self._raw_energy.values) - range_ = range(n)[self._steps] - start = range_.start + 1 - stop = range_.stop - return f"step {stop} of range {start}:{stop}" - elif self._steps == -1 or self._steps is None: - return "final step" - else: - return f"step {self._steps + 1}" - - def _default_dict(self): - raw_values = np.array(self._raw_energy.values).T - return { - convert.text_to_string(label).strip(): value[self._steps_or_last] - for label, value in zip(self._raw_energy.labels, raw_values) - } - - def _default_dict_all(self): - raw_values = np.array(self._raw_energy.values).T - return { - convert.text_to_string(label).strip(): value[:] - for label, value in zip(self._raw_energy.labels, raw_values) - } - - def _read_data(self, tree, steps): - maps = {1: self._init_selection_dict()} - selector = index.Selector(maps, self._raw_energy.values) - for selection in tree.selections(): - yield selector.label(selection), selector[selection][steps] - - def _init_selection_dict(self): - return { - selection: idx - for idx, label in enumerate(self._raw_energy.labels) - for selection in _SELECTIONS.get(convert.text_to_string(label).strip(), ()) - } - - def _make_series(self, yaxes, tree): - n = len(self._raw_energy.values) - slice_ = self._to_slice - steps = np.arange(n)[slice_] + 1 - return [ - graph.Series(x=steps, y=values, label=label, y2=yaxes.use_y2(label)) - for label, values in self._read_data(tree, slice_) - ] - - -@quantity("energy") -@documentation.format(examples=slice_.examples("energy")) -class Energy(graph.Mixin): - """The energy data for one or several steps of a relaxation or MD simulation. - - You can use this class to inspect how the ionic relaxation converges or - during an MD simulation whether the total energy is conserved. The total - energy of the system is one of the most important results to analyze materials. - Total energy differences of different atom arrangements reveal which structure - is more stable. Even when the number of atoms are different between two - systems, you may be able to compare the total energies by adding a corresponding - amount of single atom energies. In this case, you need to double check the - convergence because some error cancellation does not happen if the number of - atoms is changes. Finally, monitoring the total energy can reveal insights - about the stability of the thermostat. - - {examples} - """ - - def __init__(self, source, quantity_name: str = "energy", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_energy: raw.Energy): - """Create an Energy dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_energy)) - - def __getitem__(self, steps) -> "Energy": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw_data): - return EnergyHandler.from_data(raw_data, steps=self._steps) - - @documentation.format( - selection=_selection_string("all energies"), - examples=slice_.examples("energy", "to_dict"), - ) - def read(self, selection=None) -> dict: - """Read the energy data and store it in a dictionary. - - Parameters - ---------- - {selection} - - Returns - ------- - dict - Contains the exact labels corresponding to the selection and the - associated energies for every selected ionic step. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EnergyHandler.to_dict, - ) - - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - @documentation.format( - selection=_selection_string("the total energy"), - examples=slice_.examples("energy", "to_graph"), - ) - def to_graph(self, selection="TOTEN") -> graph.Graph: - """Read the energy data and generate a figure of the selected components. - - Parameters - ---------- - {selection} - - Returns - ------- - Graph - figure containing the selected energies for every selected ionic step. - - {examples} - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EnergyHandler.to_graph, - ) - - @documentation.format( - selection=_selection_string("the total energy"), - examples=slice_.examples("energy", "to_numpy"), - ) - def to_numpy(self, selection="TOTEN") -> np.ndarray: - """Read the energy of the selected steps. - - Parameters - ---------- - {selection} - - Returns - ------- - float or np.ndarray or tuple - Contains energies associated with the selection for the selected ionic step(s). - When only a single step is inquired, result is a float otherwise an array. - If you select multiple quantities a tuple of them is returned. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EnergyHandler.to_numpy, - ) - - def selections(self, selection: str | None = None) -> dict: - """Return a dictionary describing what kind of energies are available. - - Returns - ------- - - - Dictionary containing available selection options with their possible values. - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EnergyHandler.selections, - ) - - def __str__(self, selection: str | None = None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - EnergyHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - self._handler_factory, - EnergyHandler.to_database, - ) - - -class _YAxes: - def __init__(self, tree): - uses = set(self._is_temperature(selection) for selection in tree.selections()) - use_energy = False in uses - self.use_both = len(uses) == 2 - self.ylabel = "Energy (eV)" if use_energy else "Temperature (K)" - self.y2label = "Temperature (K)" if self.use_both else None - - def _is_temperature(self, selection): - choices = _SELECTIONS["temperature TEIN"] - return any(select.contains(selection, choice) for choice in choices) - - def use_y2(self, label): - choices = _SELECTIONS["temperature TEIN"] - return self.use_both and label in choices +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy + +import numpy as np + +from py4vasp import raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import Energy_DB +from py4vasp._third_party import graph +from py4vasp._util import convert, documentation, index, select + + +def _selection_string(default): + return f"""\ +selection : str or None + String specifying the labels of the energy to be read. If no energy is selected + this will default to selecting {default}. Separate distinct labels by commas or + whitespace. You can add or subtract different contributions e.g. `TOTEN + EKIN`. + For a complete list of all possible selections, please use + + >>> calculation.energy.selections() +""" + + +_SELECTIONS = { + "ion-electron TOTEN": ["ion_electron", "TOTEN"], + "kinetic energy EKIN": ["kinetic_energy", "EKIN"], + "kin. lattice EKIN_LAT": ["kinetic_lattice", "EKIN_LAT"], + "temperature TEIN": ["temperature", "TEIN"], + "nose potential ES": ["nose_potential", "ES"], + "nose kinetic EPS": ["nose_kinetic", "EPS"], + "total energy ETOTAL": ["total_energy", "ETOTAL"], + "free energy TOTEN": ["free_energy", "TOTEN"], + "energy without entropy": ["without_entropy", "ENOENT"], + "energy(sigma->0)": ["sigma_0", "ESIG0"], + "step STEP": ["step", "STEP"], + "One el. energy E1": ["one_electron", "E1"], + "Hartree energy -DENC": ["Hartree", "hartree", "DENC"], + "exchange EXHF": ["exchange", "EXHF"], + "free energy TOTEN": ["free_energy", "TOTEN"], + "free energy cap TOTENCAP": ["cap", "TOTENCAP"], + "weight WEIGHT": ["weight", "WEIGHT"], +} + +_DB_KEYS = { + "ion-electron TOTEN": "ion_electron", + "kinetic energy EKIN": "kinetic_energy", + "kin. lattice EKIN_LAT": "kinetic_energy_lattice", + "temperature TEIN": "temperature", + "nose potential ES": "nose_potential", + "nose kinetic EPS": "nose_kinetic", + "total energy ETOTAL": "total_energy", + "free energy TOTEN": "free_energy", + "energy without entropy": "energy_without_entropy", + "energy(sigma->0)": "energy_sigma_0", + "step STEP": "step", + "One el. energy E1": "one_electron_energy", + "Hartree energy -DENC": "hartree_energy", + "exchange EXHF": "exchange_energy", + "free energy TOTEN": "free_energy", + "free energy cap TOTENCAP": "free_energy_cap", + "weight WEIGHT": "weight", +} + + +@documentation.format(examples=slice_.examples("energy")) +class EnergyHandler: + """Handler for energy data — performs all data access and transformation logic.""" + + def __init__(self, raw_energy: raw.Energy, steps=None): + self._raw_energy = raw_energy + self._steps = steps + + @classmethod + def from_data(cls, raw_energy: raw.Energy, steps=None) -> "EnergyHandler": + return cls(raw_energy, steps) + + def __str__(self) -> str: + text = f"Energies at {self._step_string()}:" + values = self._raw_energy.values[self._last_step_in_slice] + for label, value in zip(self._raw_energy.labels, values): + label_str = f"{convert.text_to_string(label):23.23}" + text += f"\n {label_str}={value:17.6f}" + return text + + def to_dict(self, selection=None) -> dict: + if selection is None: + return self._default_dict() + tree = select.Tree.from_selection(selection) + return dict(self._read_data(tree, self._steps_or_last)) + + def to_graph(self, selection="TOTEN") -> graph.Graph: + tree = select.Tree.from_selection(selection) + yaxes = _YAxes(tree) + return graph.Graph( + series=self._make_series(yaxes, tree), + xlabel="Step", + ylabel=yaxes.ylabel, + y2label=yaxes.y2label, + ) + + def to_numpy(self, selection="TOTEN") -> np.ndarray: + tree = select.Tree.from_selection(selection) + return np.squeeze( + [values for _, values in self._read_data(tree, self._steps_or_last)] + ) + + def selections(self) -> dict: + components = list(self._init_selection_dict().keys()) + return {"energy": [], "component": components} + + def to_database(self) -> dict: + default_dict = self._default_dict_all() + energy_dict = {} + for original_label, db_key in _DB_KEYS.items(): + v = default_dict.get(original_label, None) + if v is not None: + v = np.array(v) + energy_dict[f"{db_key}_initial"] = None if v is None else float(v[0]) + if (db_key != "step") and (v is not None): + energy_dict[f"{db_key}_min"] = float(np.min(v)) + energy_dict[f"{db_key}_step_min"] = int(np.argmin(v)) + energy_dict[f"{db_key}_final"] = None if v is None else float(v[-1]) + extra_dict = {} + for k, v in default_dict.items(): + if k not in _DB_KEYS: + vs = np.array(v) if v is not None else None + key = convert.text_to_string(k).strip().lower().replace(" ", "_") + extra_dict[f"{key}_initial"] = None if vs is None else float(vs[0]) + extra_dict[f"{key}_min"] = None if vs is None else float(np.min(vs)) + extra_dict[f"{key}_step_min"] = ( + None if vs is None else int(np.argmin(vs)) + ) + extra_dict[f"{key}_final"] = None if vs is None else float(vs[-1]) + return Energy_DB(**energy_dict, other_energy_data=extra_dict) + + @property + def _steps_or_last(self): + if self._steps is None: + return -1 + return self._steps + + @property + def _last_step_in_slice(self): + if self._steps is None or self._steps == -1: + return -1 + if isinstance(self._steps, slice): + return (self._steps.stop or 0) - 1 + return self._steps + + @property + def _to_slice(self): + if self._steps is None or self._steps == -1: + return slice(-1, None) + if isinstance(self._steps, slice): + return self._steps + return slice(self._steps, self._steps + 1) + + def _step_string(self): + if isinstance(self._steps, slice): + n = len(self._raw_energy.values) + range_ = range(n)[self._steps] + start = range_.start + 1 + stop = range_.stop + return f"step {stop} of range {start}:{stop}" + elif self._steps == -1 or self._steps is None: + return "final step" + else: + return f"step {self._steps + 1}" + + def _default_dict(self): + raw_values = np.array(self._raw_energy.values).T + return { + convert.text_to_string(label).strip(): value[self._steps_or_last] + for label, value in zip(self._raw_energy.labels, raw_values) + } + + def _default_dict_all(self): + raw_values = np.array(self._raw_energy.values).T + return { + convert.text_to_string(label).strip(): value[:] + for label, value in zip(self._raw_energy.labels, raw_values) + } + + def _read_data(self, tree, steps): + maps = {1: self._init_selection_dict()} + selector = index.Selector(maps, self._raw_energy.values) + for selection in tree.selections(): + yield selector.label(selection), selector[selection][steps] + + def _init_selection_dict(self): + return { + selection: idx + for idx, label in enumerate(self._raw_energy.labels) + for selection in _SELECTIONS.get(convert.text_to_string(label).strip(), ()) + } + + def _make_series(self, yaxes, tree): + n = len(self._raw_energy.values) + slice_ = self._to_slice + steps = np.arange(n)[slice_] + 1 + return [ + graph.Series(x=steps, y=values, label=label, y2=yaxes.use_y2(label)) + for label, values in self._read_data(tree, slice_) + ] + + +@quantity("energy") +@documentation.format(examples=slice_.examples("energy")) +class Energy(graph.Mixin): + """The energy data for one or several steps of a relaxation or MD simulation. + + You can use this class to inspect how the ionic relaxation converges or + during an MD simulation whether the total energy is conserved. The total + energy of the system is one of the most important results to analyze materials. + Total energy differences of different atom arrangements reveal which structure + is more stable. Even when the number of atoms are different between two + systems, you may be able to compare the total energies by adding a corresponding + amount of single atom energies. In this case, you need to double check the + convergence because some error cancellation does not happen if the number of + atoms is changes. Finally, monitoring the total energy can reveal insights + about the stability of the thermostat. + + {examples} + """ + + def __init__(self, source, quantity_name: str = "energy", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_energy: raw.Energy): + """Create an Energy dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_energy)) + + def __getitem__(self, steps) -> "Energy": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw_data): + return EnergyHandler.from_data(raw_data, steps=self._steps) + + @documentation.format( + selection=_selection_string("all energies"), + examples=slice_.examples("energy", "to_dict"), + ) + def read(self, selection=None) -> dict: + """Read the energy data and store it in a dictionary. + + Parameters + ---------- + {selection} + + Returns + ------- + dict + Contains the exact labels corresponding to the selection and the + associated energies for every selected ionic step. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EnergyHandler.to_dict, + ) + + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + @documentation.format( + selection=_selection_string("the total energy"), + examples=slice_.examples("energy", "to_graph"), + ) + def to_graph(self, selection="TOTEN") -> graph.Graph: + """Read the energy data and generate a figure of the selected components. + + Parameters + ---------- + {selection} + + Returns + ------- + Graph + figure containing the selected energies for every selected ionic step. + + {examples} + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EnergyHandler.to_graph, + ) + + @documentation.format( + selection=_selection_string("the total energy"), + examples=slice_.examples("energy", "to_numpy"), + ) + def to_numpy(self, selection="TOTEN") -> np.ndarray: + """Read the energy of the selected steps. + + Parameters + ---------- + {selection} + + Returns + ------- + float or np.ndarray or tuple + Contains energies associated with the selection for the selected ionic step(s). + When only a single step is inquired, result is a float otherwise an array. + If you select multiple quantities a tuple of them is returned. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EnergyHandler.to_numpy, + ) + + def selections(self, selection: str | None = None) -> dict: + """Return a dictionary describing what kind of energies are available. + + Returns + ------- + - + Dictionary containing available selection options with their possible values. + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EnergyHandler.selections, + ) + + def __str__(self, selection: str | None = None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + EnergyHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + self._handler_factory, + EnergyHandler.to_database, + ) + + +class _YAxes: + def __init__(self, tree): + uses = set(self._is_temperature(selection) for selection in tree.selections()) + use_energy = False in uses + self.use_both = len(uses) == 2 + self.ylabel = "Energy (eV)" if use_energy else "Temperature (K)" + self.y2label = "Temperature (K)" if self.use_both else None + + def _is_temperature(self, selection): + choices = _SELECTIONS["temperature TEIN"] + return any(select.contains(selection, choice) for choice in choices) + + def use_y2(self, label): + choices = _SELECTIONS["temperature TEIN"] + return self.use_both and label in choices diff --git a/src/py4vasp/_calculation/exciton_eigenvector.py b/src/py4vasp/_calculation/exciton_eigenvector.py index aa14bae71..5ec9ccda6 100644 --- a/src/py4vasp/_calculation/exciton_eigenvector.py +++ b/src/py4vasp/_calculation/exciton_eigenvector.py @@ -1,143 +1,143 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import numpy as np - -from py4vasp import raw -from py4vasp._calculation._dispersion import DispersionHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import ExcitonEigenvector_DB -from py4vasp._util import check, convert - - -class ExcitonEigenvectorHandler: - """Handler for BSE exciton eigenvector data.""" - - def __init__(self, raw_exciton_eigenvector: raw.ExcitonEigenvector): - self._raw_exciton_eigenvector = raw_exciton_eigenvector - - @classmethod - def from_data( - cls, raw_exciton_eigenvector: raw.ExcitonEigenvector - ) -> "ExcitonEigenvectorHandler": - return cls(raw_exciton_eigenvector) - - def __str__(self) -> str: - shape = self._raw_exciton_eigenvector.bse_index.shape - return f"""BSE eigenvector data: - {shape[1]} k-points - {shape[3]} valence bands - {shape[2]} conduction bands""" - - def to_dict(self) -> dict: - eigenvectors = convert.to_complex(self._raw_exciton_eigenvector.eigenvectors[:]) - dispersion = self._dispersion().to_dict() - shifted_eigenvalues = ( - dispersion.pop("eigenvalues") - self._raw_exciton_eigenvector.fermi_energy - ) - return { - **dispersion, - "bands": shifted_eigenvalues, - "bse_index": self._raw_exciton_eigenvector.bse_index[:] - 1, - "eigenvectors": eigenvectors, - "fermi_energy": self._raw_exciton_eigenvector.fermi_energy, - "first_valence_band": self._raw_exciton_eigenvector.first_valence_band[:] - - 1, - "first_conduction_band": self._raw_exciton_eigenvector.first_conduction_band[ - : - ] - - 1, - } - - def to_database(self) -> dict: - num_bands_valence = None - num_bands_conduction = None - num_kpoints = None - if not check.is_none(self._raw_exciton_eigenvector.bse_index): - bse_index = self._raw_exciton_eigenvector.bse_index[:] - num_bands_conduction = np.shape(bse_index)[2] - num_bands_valence = np.shape(bse_index)[3] - num_kpoints = np.shape(bse_index)[1] - return ExcitonEigenvector_DB( - num_kpoints=num_kpoints, - num_valence_bands=num_bands_valence, - num_conduction_bands=num_bands_conduction, - ) - - def _dispersion(self) -> DispersionHandler: - return DispersionHandler.from_data(self._raw_exciton_eigenvector.dispersion) - - -@quantity("eigenvector", group="exciton") -class ExcitonEigenvector: - """BSE can compute excitonic properties of materials. - - The Bethe-Salpeter Equation (BSE) accounts for electron-hole interactions - involved in excitonic processes. For systems, where excitonic excitations - matter the BSE method is an important tool. One can visualize excitonic - contributions as so-called "fatbands" plots.""" - - def __init__(self, source, quantity_name: str = "exciton_eigenvector"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data( - cls, raw_exciton_eigenvector: raw.ExcitonEigenvector - ) -> "ExcitonEigenvector": - return cls(source=DataSource(raw_exciton_eigenvector)) - - def _handler_factory(self, raw): - return ExcitonEigenvectorHandler.from_data(raw) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ExcitonEigenvectorHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Read the data into a dictionary. - - Returns - ------- - dict - The dictionary contains the relevant k-point distances and labels as well as - the electronic band eigenvalues. To produce fatband plots, use the array - *bse_index* to access the relevant quantities of the BSE eigenvectors. Note - that the dimensions of the bse_index array are **k** points, conduction - bands, valence bands and that the conduction and valence band indices may - be offset by first_valence_band and first_conduction_band, respectively. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - ExcitonEigenvectorHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - ExcitonEigenvectorHandler.from_data, - ExcitonEigenvectorHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import numpy as np + +from py4vasp import raw +from py4vasp._calculation._dispersion import DispersionHandler +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import ExcitonEigenvector_DB +from py4vasp._util import check, convert + + +class ExcitonEigenvectorHandler: + """Handler for BSE exciton eigenvector data.""" + + def __init__(self, raw_exciton_eigenvector: raw.ExcitonEigenvector): + self._raw_exciton_eigenvector = raw_exciton_eigenvector + + @classmethod + def from_data( + cls, raw_exciton_eigenvector: raw.ExcitonEigenvector + ) -> "ExcitonEigenvectorHandler": + return cls(raw_exciton_eigenvector) + + def __str__(self) -> str: + shape = self._raw_exciton_eigenvector.bse_index.shape + return f"""BSE eigenvector data: + {shape[1]} k-points + {shape[3]} valence bands + {shape[2]} conduction bands""" + + def to_dict(self) -> dict: + eigenvectors = convert.to_complex(self._raw_exciton_eigenvector.eigenvectors[:]) + dispersion = self._dispersion().to_dict() + shifted_eigenvalues = ( + dispersion.pop("eigenvalues") - self._raw_exciton_eigenvector.fermi_energy + ) + return { + **dispersion, + "bands": shifted_eigenvalues, + "bse_index": self._raw_exciton_eigenvector.bse_index[:] - 1, + "eigenvectors": eigenvectors, + "fermi_energy": self._raw_exciton_eigenvector.fermi_energy, + "first_valence_band": self._raw_exciton_eigenvector.first_valence_band[:] + - 1, + "first_conduction_band": self._raw_exciton_eigenvector.first_conduction_band[ + : + ] + - 1, + } + + def to_database(self) -> dict: + num_bands_valence = None + num_bands_conduction = None + num_kpoints = None + if not check.is_none(self._raw_exciton_eigenvector.bse_index): + bse_index = self._raw_exciton_eigenvector.bse_index[:] + num_bands_conduction = np.shape(bse_index)[2] + num_bands_valence = np.shape(bse_index)[3] + num_kpoints = np.shape(bse_index)[1] + return ExcitonEigenvector_DB( + num_kpoints=num_kpoints, + num_valence_bands=num_bands_valence, + num_conduction_bands=num_bands_conduction, + ) + + def _dispersion(self) -> DispersionHandler: + return DispersionHandler.from_data(self._raw_exciton_eigenvector.dispersion) + + +@quantity("eigenvector", group="exciton") +class ExcitonEigenvector: + """BSE can compute excitonic properties of materials. + + The Bethe-Salpeter Equation (BSE) accounts for electron-hole interactions + involved in excitonic processes. For systems, where excitonic excitations + matter the BSE method is an important tool. One can visualize excitonic + contributions as so-called "fatbands" plots.""" + + def __init__(self, source, quantity_name: str = "exciton_eigenvector"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data( + cls, raw_exciton_eigenvector: raw.ExcitonEigenvector + ) -> "ExcitonEigenvector": + return cls(source=DataSource(raw_exciton_eigenvector)) + + def _handler_factory(self, raw): + return ExcitonEigenvectorHandler.from_data(raw) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ExcitonEigenvectorHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Read the data into a dictionary. + + Returns + ------- + dict + The dictionary contains the relevant k-point distances and labels as well as + the electronic band eigenvalues. To produce fatband plots, use the array + *bse_index* to access the relevant quantities of the BSE eigenvectors. Note + that the dimensions of the bse_index array are **k** points, conduction + bands, valence bands and that the conduction and valence band indices may + be offset by first_valence_band and first_conduction_band, respectively. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + ExcitonEigenvectorHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + ExcitonEigenvectorHandler.from_data, + ExcitonEigenvectorHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/force.py b/src/py4vasp/_calculation/force.py index 31eb4cf53..a023231cf 100644 --- a/src/py4vasp/_calculation/force.py +++ b/src/py4vasp/_calculation/force.py @@ -1,352 +1,352 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy - -import numpy as np - -from py4vasp import _config, exception, raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, - slice_steps, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw.data_db import Force_DB -from py4vasp._third_party import view -from py4vasp._util import check - - -class ForceHandler: - """Handler for force data — performs all data access and transformation logic.""" - - force_rescale = 1.5 - "Scaling constant to convert forces to Å." - - def __init__(self, raw_force: raw.Force, steps=None): - self._raw_force = raw_force - self._steps = steps - - @classmethod - def from_data(cls, raw_force: raw.Force, steps=None) -> "ForceHandler": - return cls(raw_force, steps=steps) - - def __str__(self) -> str: - result = """ -POSITION TOTAL-FORCE (eV/Angst) ------------------------------------------------------------------------------------ - """.strip() - step = self._last_step - structure = StructureHandler.from_data(self._raw_force.structure, steps=step) - positions = structure.cartesian_positions() - forces = np.array(self._raw_force.forces)[step] - position_to_string = lambda pos: " ".join(f"{x:12.5f}" for x in pos) - force_to_string = lambda f: " ".join(f"{x:13.6f}" for x in f) - for position, force in zip(positions, forces): - result += f"\n{position_to_string(position)} {force_to_string(force)}" - return result - - def to_dict(self) -> dict: - """Read the forces into a dictionary. - - Forces and associated structural information for one or more selected steps of - the trajectory are returned in a dictionary. This includes the lattice vectors, - atomic positions, and atomic species in addition to the forces acting on each atom. - The forces are in Cartesian coordinates and in units of eV/Å. - - Returns - ------- - dict - Contains the forces for all selected steps and the structural information - to know on which atoms the forces act. - """ - structure = StructureHandler.from_data( - self._raw_force.structure, steps=self._steps - ) - return { - "structure": structure.to_dict(), - "forces": slice_steps( - np.array(self._raw_force.forces), self._steps, default_ndim=2 - ), - } - - def to_database(self) -> dict: - """Serialize force statistics to the database format.""" - if check.is_none(self._raw_force.forces): - raise exception.NoData("No force data available to write to database.") - forces = np.array(self._raw_force.forces) - if forces.ndim == 2: - final_force_norms = np.linalg.norm(forces, axis=-1) - initial_force_norms = final_force_norms.copy() - else: - final_force_norms = np.linalg.norm(forces[-1], axis=-1) - initial_force_norms = np.linalg.norm(forces[0], axis=-1) - return Force_DB( - final_force_min=float(np.min(final_force_norms)), - final_force_median=float(np.median(final_force_norms)), - final_force_mean=float(np.mean(final_force_norms)), - final_force_max=float(np.max(final_force_norms)), - final_index_force_max=int(np.argmax(final_force_norms)), - initial_force_min=float(np.min(initial_force_norms)), - initial_force_max=float(np.max(initial_force_norms)), - initial_index_force_max=int(np.argmax(initial_force_norms)), - ) - - def to_view(self, supercell=None): - """Visualize the forces showing arrows at the atoms.""" - structure = StructureHandler.from_data(self._raw_force.structure) - viewer = structure.to_view(supercell) - forces = self.force_rescale * slice_steps( - np.array(self._raw_force.forces), self._steps, default_ndim=2 - ) - if forces.ndim == 2: - forces = forces[np.newaxis] - ion_arrow = view.IonArrow( - quantity=forces, - label="forces", - color=_config.VASP_COLORS["purple"], - radius=0.2, - ) - viewer.ion_arrows = [ion_arrow] - return viewer - - def number_steps(self) -> int: - """Return the number of forces in the trajectory.""" - n = len(np.array(self._raw_force.forces)) - return len(range(n)[self._to_slice]) - - @property - def _last_step(self): - if self._steps is None or self._steps == -1: - return -1 - if isinstance(self._steps, slice): - stop = self._steps.stop - return (stop - 1) if stop is not None else -1 - return self._steps - - @property - def _to_slice(self): - if self._steps is None or self._steps == -1: - return slice(-1, None) - if isinstance(self._steps, slice): - return self._steps - return slice(self._steps, self._steps + 1) - - -@quantity("force") -class Force(view.Mixin): - """The forces determine the path of the atoms in a trajectory. - - You can use this class to analyze the forces acting on the atoms. The forces - are the first derivative of the DFT total energy. The forces being small is - an important criterion for the convergence of a relaxation calculation. The - size of the forces is also related to the maximal time step in MD simulations. - When you choose a too large time step, the forces become large and the atoms - may move too much in a single step leading to an unstable trajectory. You can - use this class to visualize the forces in a trajectory or read the values to - analyze them numerically. - - Examples - -------- - Let us create some example data so that we can illustrate how to use this class. - Of course you can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you access the forces, the result will depend on the steps that you selected - with the [] operator. Without any selection the results from the final step will be - used. - - >>> calculation.force.number_steps() - 1 - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.force[:].number_steps() - 4 - - You can also select specific steps or a subset of steps as follows - - >>> calculation.force[3].number_steps() - 1 - >>> calculation.force[1:4].number_steps() - 3 - """ - - force_rescale = ForceHandler.force_rescale - - def __init__(self, source, quantity_name: str = "force", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_force: raw.Force) -> "Force": - """Create a Force dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_force)) - - def __getitem__(self, steps) -> "Force": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw): - return ForceHandler.from_data(raw, steps=self._steps) - - def __str__(self, selection=None) -> str: - "Convert the forces to a format similar to the OUTCAR file." - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ForceHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def read(self) -> dict: - """Read the forces into a dictionary. - - Forces and associated structural information for one or more selected steps of - the trajectory are returned in a dictionary. This includes the lattice vectors, - atomic positions, and atomic species in addition to the forces acting on each atom. - The forces are in Cartesian coordinates and in units of eV/Å. - - Returns - ------- - dict - Contains the forces for all selected steps and the structural information - to know on which atoms the forces act. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `read` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. The structure is included to provide the necessary context for - the forces. - - >>> calculation.force.read() - {'structure': {...}, 'forces': array([[...]])} - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the forces contain an additional dimension for the - different steps. - - >>> calculation.force[:].read() - {'structure': {...}, 'forces': array([[[...]]])} - - You can also select specific steps or a subset of steps as follows - - >>> calculation.force[1].read() - {'structure': {...}, 'forces': array([[...]])} - >>> calculation.force[0:2].read() - {'structure': {...}, 'forces': array([[[...]]])} - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - ForceHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def to_view(self, supercell=None) -> view.View: - """Visualize the forces showing arrows at the atoms. - - This method adds arrows to the atoms in the structure sized according to the - strength of the force. The length of the arrows is scaled by a constant - `force_rescale` to convert from eV/Å to a length in Å. - - Parameters - ---------- - supercell : int or np.ndarray - If present the structure is replicated the specified number of times - along each direction. - - Returns - ------- - View - Shows the structure with cell and all atoms adding arrows to the atoms - sized according to the strength of the force. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_view` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.force.to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='forces', ...)], ...) - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.force[:].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='forces', ...)], ...) - - You can also select specific steps or a subset of steps as follows - - >>> calculation.force[1].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='forces', ...)], ...) - >>> calculation.force[0:2].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...], [...]]]), label='forces', ...)], ...) - - You may also replicate the structure by specifying a supercell. - - >>> calculation.force.to_view(supercell=2) - View(..., supercell=array([2, 2, 2]), ...) - - The supercell size can also be different for the different directions. - - >>> calculation.force.to_view(supercell=[2,3,1]) - View(..., supercell=array([2, 3, 1]), ...) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - ForceHandler.to_view, - supercell, - ) - - def number_steps(self) -> int: - """Return the number of forces in the trajectory.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - ForceHandler.number_steps, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - ForceHandler.from_data, - ForceHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy + +import numpy as np + +from py4vasp import _config, exception, raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, + slice_steps, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw.data_db import Force_DB +from py4vasp._third_party import view +from py4vasp._util import check + + +class ForceHandler: + """Handler for force data — performs all data access and transformation logic.""" + + force_rescale = 1.5 + "Scaling constant to convert forces to Å." + + def __init__(self, raw_force: raw.Force, steps=None): + self._raw_force = raw_force + self._steps = steps + + @classmethod + def from_data(cls, raw_force: raw.Force, steps=None) -> "ForceHandler": + return cls(raw_force, steps=steps) + + def __str__(self) -> str: + result = """ +POSITION TOTAL-FORCE (eV/Angst) +----------------------------------------------------------------------------------- + """.strip() + step = self._last_step + structure = StructureHandler.from_data(self._raw_force.structure, steps=step) + positions = structure.cartesian_positions() + forces = np.array(self._raw_force.forces)[step] + position_to_string = lambda pos: " ".join(f"{x:12.5f}" for x in pos) + force_to_string = lambda f: " ".join(f"{x:13.6f}" for x in f) + for position, force in zip(positions, forces): + result += f"\n{position_to_string(position)} {force_to_string(force)}" + return result + + def to_dict(self) -> dict: + """Read the forces into a dictionary. + + Forces and associated structural information for one or more selected steps of + the trajectory are returned in a dictionary. This includes the lattice vectors, + atomic positions, and atomic species in addition to the forces acting on each atom. + The forces are in Cartesian coordinates and in units of eV/Å. + + Returns + ------- + dict + Contains the forces for all selected steps and the structural information + to know on which atoms the forces act. + """ + structure = StructureHandler.from_data( + self._raw_force.structure, steps=self._steps + ) + return { + "structure": structure.to_dict(), + "forces": slice_steps( + np.array(self._raw_force.forces), self._steps, default_ndim=2 + ), + } + + def to_database(self) -> dict: + """Serialize force statistics to the database format.""" + if check.is_none(self._raw_force.forces): + raise exception.NoData("No force data available to write to database.") + forces = np.array(self._raw_force.forces) + if forces.ndim == 2: + final_force_norms = np.linalg.norm(forces, axis=-1) + initial_force_norms = final_force_norms.copy() + else: + final_force_norms = np.linalg.norm(forces[-1], axis=-1) + initial_force_norms = np.linalg.norm(forces[0], axis=-1) + return Force_DB( + final_force_min=float(np.min(final_force_norms)), + final_force_median=float(np.median(final_force_norms)), + final_force_mean=float(np.mean(final_force_norms)), + final_force_max=float(np.max(final_force_norms)), + final_index_force_max=int(np.argmax(final_force_norms)), + initial_force_min=float(np.min(initial_force_norms)), + initial_force_max=float(np.max(initial_force_norms)), + initial_index_force_max=int(np.argmax(initial_force_norms)), + ) + + def to_view(self, supercell=None): + """Visualize the forces showing arrows at the atoms.""" + structure = StructureHandler.from_data(self._raw_force.structure) + viewer = structure.to_view(supercell) + forces = self.force_rescale * slice_steps( + np.array(self._raw_force.forces), self._steps, default_ndim=2 + ) + if forces.ndim == 2: + forces = forces[np.newaxis] + ion_arrow = view.IonArrow( + quantity=forces, + label="forces", + color=_config.VASP_COLORS["purple"], + radius=0.2, + ) + viewer.ion_arrows = [ion_arrow] + return viewer + + def number_steps(self) -> int: + """Return the number of forces in the trajectory.""" + n = len(np.array(self._raw_force.forces)) + return len(range(n)[self._to_slice]) + + @property + def _last_step(self): + if self._steps is None or self._steps == -1: + return -1 + if isinstance(self._steps, slice): + stop = self._steps.stop + return (stop - 1) if stop is not None else -1 + return self._steps + + @property + def _to_slice(self): + if self._steps is None or self._steps == -1: + return slice(-1, None) + if isinstance(self._steps, slice): + return self._steps + return slice(self._steps, self._steps + 1) + + +@quantity("force") +class Force(view.Mixin): + """The forces determine the path of the atoms in a trajectory. + + You can use this class to analyze the forces acting on the atoms. The forces + are the first derivative of the DFT total energy. The forces being small is + an important criterion for the convergence of a relaxation calculation. The + size of the forces is also related to the maximal time step in MD simulations. + When you choose a too large time step, the forces become large and the atoms + may move too much in a single step leading to an unstable trajectory. You can + use this class to visualize the forces in a trajectory or read the values to + analyze them numerically. + + Examples + -------- + Let us create some example data so that we can illustrate how to use this class. + Of course you can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you access the forces, the result will depend on the steps that you selected + with the [] operator. Without any selection the results from the final step will be + used. + + >>> calculation.force.number_steps() + 1 + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.force[:].number_steps() + 4 + + You can also select specific steps or a subset of steps as follows + + >>> calculation.force[3].number_steps() + 1 + >>> calculation.force[1:4].number_steps() + 3 + """ + + force_rescale = ForceHandler.force_rescale + + def __init__(self, source, quantity_name: str = "force", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_force: raw.Force) -> "Force": + """Create a Force dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_force)) + + def __getitem__(self, steps) -> "Force": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw): + return ForceHandler.from_data(raw, steps=self._steps) + + def __str__(self, selection=None) -> str: + "Convert the forces to a format similar to the OUTCAR file." + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ForceHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def read(self) -> dict: + """Read the forces into a dictionary. + + Forces and associated structural information for one or more selected steps of + the trajectory are returned in a dictionary. This includes the lattice vectors, + atomic positions, and atomic species in addition to the forces acting on each atom. + The forces are in Cartesian coordinates and in units of eV/Å. + + Returns + ------- + dict + Contains the forces for all selected steps and the structural information + to know on which atoms the forces act. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `read` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. The structure is included to provide the necessary context for + the forces. + + >>> calculation.force.read() + {'structure': {...}, 'forces': array([[...]])} + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the forces contain an additional dimension for the + different steps. + + >>> calculation.force[:].read() + {'structure': {...}, 'forces': array([[[...]]])} + + You can also select specific steps or a subset of steps as follows + + >>> calculation.force[1].read() + {'structure': {...}, 'forces': array([[...]])} + >>> calculation.force[0:2].read() + {'structure': {...}, 'forces': array([[[...]]])} + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + ForceHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def to_view(self, supercell=None) -> view.View: + """Visualize the forces showing arrows at the atoms. + + This method adds arrows to the atoms in the structure sized according to the + strength of the force. The length of the arrows is scaled by a constant + `force_rescale` to convert from eV/Å to a length in Å. + + Parameters + ---------- + supercell : int or np.ndarray + If present the structure is replicated the specified number of times + along each direction. + + Returns + ------- + View + Shows the structure with cell and all atoms adding arrows to the atoms + sized according to the strength of the force. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_view` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.force.to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='forces', ...)], ...) + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.force[:].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='forces', ...)], ...) + + You can also select specific steps or a subset of steps as follows + + >>> calculation.force[1].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='forces', ...)], ...) + >>> calculation.force[0:2].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...], [...]]]), label='forces', ...)], ...) + + You may also replicate the structure by specifying a supercell. + + >>> calculation.force.to_view(supercell=2) + View(..., supercell=array([2, 2, 2]), ...) + + The supercell size can also be different for the different directions. + + >>> calculation.force.to_view(supercell=[2,3,1]) + View(..., supercell=array([2, 3, 1]), ...) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + ForceHandler.to_view, + supercell, + ) + + def number_steps(self) -> int: + """Return the number of forces in the trajectory.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + ForceHandler.number_steps, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + ForceHandler.from_data, + ForceHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/kpoint.py b/src/py4vasp/_calculation/kpoint.py index 724c2e42d..3df6a535a 100644 --- a/src/py4vasp/_calculation/kpoint.py +++ b/src/py4vasp/_calculation/kpoint.py @@ -1,541 +1,541 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import functools -from contextlib import suppress -from fractions import Fraction -from typing import Any - -import numpy as np -from numpy.typing import ArrayLike - -from py4vasp import exception -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Kpoint_DB -from py4vasp._util import check, convert - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - exception.RefinementError, -) - - -def _safe_call(func): - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - return func() - return None - - -class KpointHandler: - """Handler for k-point data.""" - - def __init__(self, raw_kpoint: raw.Kpoint): - self._raw_kpoint = raw_kpoint - - @classmethod - def from_data(cls, raw_kpoint: raw.Kpoint) -> "KpointHandler": - return cls(raw_kpoint) - - def __str__(self): - text = f"""k-points -{len(self._raw_kpoint.coordinates)} -reciprocal""" - for kpoint, weight in zip( - self._raw_kpoint.coordinates, self._raw_kpoint.weights - ): - text += "\n" + f"{kpoint[0]} {kpoint[1]} {kpoint[2]} {weight}" - return text - - def to_dict(self) -> dict[str, Any]: - labels = self.labels() - labels_dict = {} if labels is None else {"labels": labels} - return { - "mode": self.mode(), - "line_length": self.line_length(), - "number_kpoints": self.number_kpoints(), - "coordinates": self._raw_kpoint.coordinates[:], - "weights": self._raw_kpoint.weights[:], - **labels_dict, - } - - def to_database(self) -> dict: - number_x = self._raw_kpoint.number_x - number_y = self._raw_kpoint.number_y - number_z = self._raw_kpoint.number_z - has_grid = not (any(check.is_none(n) for n in (number_x, number_y, number_z))) - grid_kpoints = None if not has_grid else [number_x, number_y, number_z] - user_labels = None - if not check.is_none(self._raw_kpoint.label_indices): - user_labels = [k for k in self._labels_from_file() if k != ""] - user_labels = None if len(user_labels) == 0 else user_labels - sampled_points = sorted(set(user_labels)) if user_labels is not None else None - mode = _safe_call(self.mode) - line_length = _safe_call(self.line_length) - num_kpoints_total = _safe_call(self.number_kpoints) - num_lines = _safe_call(self.number_lines) - return Kpoint_DB( - mode=mode, - line_length=line_length, - num_kpoints_total=num_kpoints_total, - num_lines=num_lines, - num_kpoints_grid=grid_kpoints, - labels=user_labels, - labels_unique=sampled_points, - ) - - def line_length(self) -> int: - if self.mode() == "line": - return self._raw_kpoint.number - return self.number_kpoints() - - def number_lines(self) -> int: - return int(self.number_kpoints() // self.line_length()) - - def number_kpoints(self) -> int: - return len(self._raw_kpoint.coordinates) - - def distances(self) -> np.ndarray: - cell = _last_step(self._raw_kpoint.cell.lattice_vectors) - cartesian_kpoints = np.linalg.solve(cell, self._raw_kpoint.coordinates[:].T).T - kpoint_lines = np.split(cartesian_kpoints, self.number_lines()) - kpoint_norms = [_line_distances(line) for line in kpoint_lines] - concatenate_distances = lambda current, addition: ( - np.concatenate((current, addition + current[-1])) - ) - return functools.reduce(concatenate_distances, kpoint_norms) - - def mode(self) -> str: - mode = convert.text_to_string(self._raw_kpoint.mode).strip() or "# empty string" - first_char = mode[0].lower() - if first_char == "a": - return "automatic" - elif first_char == "b": - return "generating lattice" - elif first_char == "e": - return "explicit" - elif first_char == "g": - return "gamma" - elif first_char == "l": - return "line" - elif first_char == "m": - return "monkhorst" - else: - raise exception.RefinementError( - f"Could not understand the mode '{mode}' when refining the raw kpoints data." - ) - - def labels(self) -> list[str] | None: - if not self._raw_kpoint.label_indices.is_none(): - return self._labels_from_file() - elif self.mode() == "line": - return self._labels_at_band_edges() - else: - return None - - def path_indices(self, start: ArrayLike, finish: ArrayLike) -> np.ndarray: - direction = np.array(finish) - np.array(start) - deltas = self._raw_kpoint.coordinates - np.array(start) - areas = np.linalg.norm(np.cross(direction, deltas), axis=1) - return np.flatnonzero(np.isclose(areas, 0)) - - def _labels_from_file(self): - labels = [""] * len(self._raw_kpoint.coordinates) - for label, index in zip(self._raw_kpoint.labels, self._raw_indices()): - labels[index] = convert.text_to_string(label.strip()) - return labels - - def _raw_indices(self): - indices = np.array(self._raw_kpoint.label_indices) - if self.mode() == "line": - line_length = self.line_length() - return line_length * (indices // 2) - (indices + 1) % 2 - else: - return indices - 1 - - def _labels_at_band_edges(self): - line_length = self.line_length() - band_edge = lambda index: not (0 < index % line_length < line_length - 1) - return [ - _kpoint_label(kpoint) if band_edge(index) else "" - for index, kpoint in enumerate(self._raw_kpoint.coordinates) - ] - - def _reciprocal_lattice_vectors(self): - scale = self._raw_kpoint.cell.scale - lattice_vectors = scale * _last_step(self._raw_kpoint.cell.lattice_vectors) - volume = np.linalg.det(lattice_vectors) - return (2.0 * np.pi / volume) * np.array( - [ - np.cross(lattice_vectors[1], lattice_vectors[2]), - np.cross(lattice_vectors[2], lattice_vectors[0]), - np.cross(lattice_vectors[0], lattice_vectors[1]), - ] - ) - - -@quantity("kpoint") -class Kpoint: - """The **k**-point mesh used in the VASP calculation. - - In VASP calculations, **k** points play an important role in discretizing the - Brillouin zone of a crystal. For self-consistent DFT calculations, typically a - regular grid of **k** points is employed to sample the Brillouin zone. A - sufficiently dense **k**-points mesh is critical for the precision of your DFT - calculation, so make sure to test the results for different meshes. Denser - **k** point meshes provide more accurate results but also demand greater - computational resources. - - Another common use case is irregular meshes in non-self-consistent calculations. - In particular in band structure analysis, one employs a mesh along specific lines - in the Brillouin zone. The line mode involves connecting high-symmetry points and - calculating the electronic band structure along these paths. - - This class provides utility functionality to extract information about either of - the aforementioned use cases. As such it is mostly used as a helper class for - other postprocessing classes to extract the required information, e.g., to - generate a band structure. It may also be used to programmatically analyze the - selected **k** point mesh or take subsets along high symmetry lines. - """ - - def __init__(self, source, quantity_name="kpoint"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_kpoint): - return cls(source=DataSource(raw_kpoint)) - - def _handler_factory(self, raw): - return KpointHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self, selection=None) -> dict: - """Read the **k** points data into a dictionary. - - Parameters - ---------- - selection : str, optional - You can select "kpoints_opt" or "kpoints_wan" here, to read from those - meshes instead of the default one defined by the KPOINTS file. - - Returns - ------- - - - Contains the coordinates of the **k** points (in crystal units) as - well as their weights used for integrations. Moreover, some data - specified in the input file of Vasp are transferred such as the mode - used to generate the **k** points, the line length (if line mode was - used), and any labels set for specific points. - - Examples - -------- - Read the **k** points data into a dictionary: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> calculation.kpoint.read() - {'mode': ..., 'line_length': ..., 'number_kpoints': ..., 'coordinates': array(...), 'weights': array(...)} - - Select the **k** points from the "kpoints_opt" mesh instead of the default one: - - >>> calculation.kpoint.read(selection="kpoints_opt") - {'mode': ..., 'line_length': ..., 'number_kpoints': ..., 'coordinates': array(...), 'weights': array(...), 'labels': ...} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.to_dict, - ) - - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - def line_length(self, selection=None) -> int: - """Get the number of points per line in the Brillouin zone. - - Returns - ------- - - - The number of points used to sample a single line. - - Examples - -------- - Get the number of points per line in the Brillouin zone: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> calculation.kpoint.line_length() - 48 - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.line_length, - ) - - def number_lines(self, selection=None) -> int: - """Get the number of lines in the Brillouin zone. - - Returns - ------- - - - The number of lines the band structure contains. For regular meshes this is - set to 1. - - Examples - -------- - Get the number of lines in the Brillouin zone for the "kpoints_opt" mesh: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> calculation.kpoint.number_lines(selection="kpoints_opt") - 4 - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.number_lines, - ) - - def number_kpoints(self, selection=None) -> int: - """Get the number of points in the Brillouin zone. - - Returns - ------- - - - The number of points used to sample the Brillouin zone. - - Examples - -------- - Get the number of points in the Brillouin zone: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> calculation.kpoint.number_kpoints() - 48 - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.number_kpoints, - ) - - def distances(self, selection=None) -> np.ndarray: - """Convert the coordinates of the **k** points into a one dimensional array. - - For every line in the Brillouin zone, the distance between each **k** point - and the start of the line is calculated. Then the distances of different - lines are concatenated into a single list. This routine is mostly useful - to plot data along high-symmetry lines like band structures. - - Returns - ------- - - - A reduction of the **k** points onto a one-dimensional array based - on the distance between the points. - - Examples - -------- - Convert the coordinates of the **k** points into a one dimensional array: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> calculation.kpoint.distances() - array([...]) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.distances, - ) - - def mode(self, selection=None) -> str: - """Get the **k**-point generation mode specified in the Vasp input file. - - Returns - ------- - - - A string representing which mode was used to setup the k-points. - - Examples - -------- - Get the **k**-point generation mode specified in the KPOINTS_OPT file: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> calculation.kpoint.mode(selection="kpoints_opt") - 'line' - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.mode, - ) - - def labels(self, selection=None) -> list[str] | None: - """Get any labels given in the input file for specific **k** points. - - The returned labels depend on the **k**-point mode and whether the user - provided explicit labels in the input file: - - - If labels are specified in the KPOINTS file, a list of strings is returned - with one entry per **k** point. Points with no user-defined label get an - empty string, while labeled points carry the name from the input file. - - If line mode is used but no labels were given, VASP automatically assigns - labels at the band edges. Interior points along the line receive an empty - string. Labels are formatted as LaTeX fractions. - - For any other mode without explicit labels (e.g., a regular Gamma or - Monkhorst-Pack grid), ``None`` is returned. - - Returns - ------- - - - A list of strings (one per **k** point) or ``None`` if no labeling is - applicable. - - Examples - -------- - If no labels were given and line mode is not used, returns None: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> result = calculation.kpoint.labels() - >>> assert result is None - - If line mode is used VASP automatically assigns labels to the band edges. In this case, - the method returns a list where band-edge points carry LaTeX-formatted coordinates and - interior points are empty strings. The example below uses the KPOINTS_OPT file: - - >>> calculation.kpoint.labels(selection="kpoints_opt") - ['$[0 0 0]$', ...] - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.labels, - ) - - def path_indices( - self, start: ArrayLike, finish: ArrayLike, selection=None - ) -> np.ndarray: - """Find linear dependent k points between start and finish. - - Loop over all possible k points and return the indices of the ones for which - k-point - start is linear dependent on finish - start. - - The primary use case is to extract a band-structure-like slice from a - regular **k**-point grid. In certain calculation types — such as - time-dependent DFT (TDDFT), GW, or BSE — VASP only supports uniform - Gamma or Monkhorst-Pack meshes. You can use this method to select all grid - points that happen to lie on a high-symmetry path. - - Parameters - ---------- - start - The starting **k** point of the path segment in fractional (crystal) - coordinates. Expects exactly 3 coordinates. - finish - The ending **k** point of the path segment in fractional (crystal) - coordinates. Expects exactly 3 coordinates. - - Returns - ------- - - - An integer array of indices (into the full **k**-point list) of all - **k** points that lie on the line segment from ``start`` to ``finish``. - - Examples - -------- - Extract all **k** points on a line through the Brillouin zone: - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - >>> start = [0, 0, 0.125] - >>> finish = [1, 0, 0.125] - >>> calculation.kpoint.path_indices(start, finish) - array([...]) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler.path_indices, - start, - finish, - ) - - def selections(self): - from py4vasp._raw import definition as raw_module - - return {self._quantity_name: list(raw_module.selections(self._quantity_name))} - - def _reciprocal_lattice_vectors(self, selection=None): - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - KpointHandler._reciprocal_lattice_vectors, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - KpointHandler.from_data, - KpointHandler.to_database, - ) - - -def _last_step(lattice_vectors): - if lattice_vectors.ndim == 2: - return lattice_vectors - else: - return lattice_vectors[-1] - - -def _line_distances(coordinates): - distances = np.zeros(len(coordinates)) - norms = np.linalg.norm(coordinates[1:] - coordinates[:-1], axis=1) - distances[1:] = np.cumsum(norms) - return distances - - -def _kpoint_label(kpoint): - fractions = [convert.Fraction(coordinate).latex() for coordinate in kpoint] - return f"$[{fractions[0]} {fractions[1]} {fractions[2]}]$" +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import functools +from contextlib import suppress +from fractions import Fraction +from typing import Any + +import numpy as np +from numpy.typing import ArrayLike + +from py4vasp import exception +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Kpoint_DB +from py4vasp._util import check, convert + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + exception.RefinementError, +) + + +def _safe_call(func): + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + return func() + return None + + +class KpointHandler: + """Handler for k-point data.""" + + def __init__(self, raw_kpoint: raw.Kpoint): + self._raw_kpoint = raw_kpoint + + @classmethod + def from_data(cls, raw_kpoint: raw.Kpoint) -> "KpointHandler": + return cls(raw_kpoint) + + def __str__(self): + text = f"""k-points +{len(self._raw_kpoint.coordinates)} +reciprocal""" + for kpoint, weight in zip( + self._raw_kpoint.coordinates, self._raw_kpoint.weights + ): + text += "\n" + f"{kpoint[0]} {kpoint[1]} {kpoint[2]} {weight}" + return text + + def to_dict(self) -> dict[str, Any]: + labels = self.labels() + labels_dict = {} if labels is None else {"labels": labels} + return { + "mode": self.mode(), + "line_length": self.line_length(), + "number_kpoints": self.number_kpoints(), + "coordinates": self._raw_kpoint.coordinates[:], + "weights": self._raw_kpoint.weights[:], + **labels_dict, + } + + def to_database(self) -> dict: + number_x = self._raw_kpoint.number_x + number_y = self._raw_kpoint.number_y + number_z = self._raw_kpoint.number_z + has_grid = not (any(check.is_none(n) for n in (number_x, number_y, number_z))) + grid_kpoints = None if not has_grid else [number_x, number_y, number_z] + user_labels = None + if not check.is_none(self._raw_kpoint.label_indices): + user_labels = [k for k in self._labels_from_file() if k != ""] + user_labels = None if len(user_labels) == 0 else user_labels + sampled_points = sorted(set(user_labels)) if user_labels is not None else None + mode = _safe_call(self.mode) + line_length = _safe_call(self.line_length) + num_kpoints_total = _safe_call(self.number_kpoints) + num_lines = _safe_call(self.number_lines) + return Kpoint_DB( + mode=mode, + line_length=line_length, + num_kpoints_total=num_kpoints_total, + num_lines=num_lines, + num_kpoints_grid=grid_kpoints, + labels=user_labels, + labels_unique=sampled_points, + ) + + def line_length(self) -> int: + if self.mode() == "line": + return self._raw_kpoint.number + return self.number_kpoints() + + def number_lines(self) -> int: + return int(self.number_kpoints() // self.line_length()) + + def number_kpoints(self) -> int: + return len(self._raw_kpoint.coordinates) + + def distances(self) -> np.ndarray: + cell = _last_step(self._raw_kpoint.cell.lattice_vectors) + cartesian_kpoints = np.linalg.solve(cell, self._raw_kpoint.coordinates[:].T).T + kpoint_lines = np.split(cartesian_kpoints, self.number_lines()) + kpoint_norms = [_line_distances(line) for line in kpoint_lines] + concatenate_distances = lambda current, addition: ( + np.concatenate((current, addition + current[-1])) + ) + return functools.reduce(concatenate_distances, kpoint_norms) + + def mode(self) -> str: + mode = convert.text_to_string(self._raw_kpoint.mode).strip() or "# empty string" + first_char = mode[0].lower() + if first_char == "a": + return "automatic" + elif first_char == "b": + return "generating lattice" + elif first_char == "e": + return "explicit" + elif first_char == "g": + return "gamma" + elif first_char == "l": + return "line" + elif first_char == "m": + return "monkhorst" + else: + raise exception.RefinementError( + f"Could not understand the mode '{mode}' when refining the raw kpoints data." + ) + + def labels(self) -> list[str] | None: + if not self._raw_kpoint.label_indices.is_none(): + return self._labels_from_file() + elif self.mode() == "line": + return self._labels_at_band_edges() + else: + return None + + def path_indices(self, start: ArrayLike, finish: ArrayLike) -> np.ndarray: + direction = np.array(finish) - np.array(start) + deltas = self._raw_kpoint.coordinates - np.array(start) + areas = np.linalg.norm(np.cross(direction, deltas), axis=1) + return np.flatnonzero(np.isclose(areas, 0)) + + def _labels_from_file(self): + labels = [""] * len(self._raw_kpoint.coordinates) + for label, index in zip(self._raw_kpoint.labels, self._raw_indices()): + labels[index] = convert.text_to_string(label.strip()) + return labels + + def _raw_indices(self): + indices = np.array(self._raw_kpoint.label_indices) + if self.mode() == "line": + line_length = self.line_length() + return line_length * (indices // 2) - (indices + 1) % 2 + else: + return indices - 1 + + def _labels_at_band_edges(self): + line_length = self.line_length() + band_edge = lambda index: not (0 < index % line_length < line_length - 1) + return [ + _kpoint_label(kpoint) if band_edge(index) else "" + for index, kpoint in enumerate(self._raw_kpoint.coordinates) + ] + + def _reciprocal_lattice_vectors(self): + scale = self._raw_kpoint.cell.scale + lattice_vectors = scale * _last_step(self._raw_kpoint.cell.lattice_vectors) + volume = np.linalg.det(lattice_vectors) + return (2.0 * np.pi / volume) * np.array( + [ + np.cross(lattice_vectors[1], lattice_vectors[2]), + np.cross(lattice_vectors[2], lattice_vectors[0]), + np.cross(lattice_vectors[0], lattice_vectors[1]), + ] + ) + + +@quantity("kpoint") +class Kpoint: + """The **k**-point mesh used in the VASP calculation. + + In VASP calculations, **k** points play an important role in discretizing the + Brillouin zone of a crystal. For self-consistent DFT calculations, typically a + regular grid of **k** points is employed to sample the Brillouin zone. A + sufficiently dense **k**-points mesh is critical for the precision of your DFT + calculation, so make sure to test the results for different meshes. Denser + **k** point meshes provide more accurate results but also demand greater + computational resources. + + Another common use case is irregular meshes in non-self-consistent calculations. + In particular in band structure analysis, one employs a mesh along specific lines + in the Brillouin zone. The line mode involves connecting high-symmetry points and + calculating the electronic band structure along these paths. + + This class provides utility functionality to extract information about either of + the aforementioned use cases. As such it is mostly used as a helper class for + other postprocessing classes to extract the required information, e.g., to + generate a band structure. It may also be used to programmatically analyze the + selected **k** point mesh or take subsets along high symmetry lines. + """ + + def __init__(self, source, quantity_name="kpoint"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_kpoint): + return cls(source=DataSource(raw_kpoint)) + + def _handler_factory(self, raw): + return KpointHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self, selection=None) -> dict: + """Read the **k** points data into a dictionary. + + Parameters + ---------- + selection : str, optional + You can select "kpoints_opt" or "kpoints_wan" here, to read from those + meshes instead of the default one defined by the KPOINTS file. + + Returns + ------- + - + Contains the coordinates of the **k** points (in crystal units) as + well as their weights used for integrations. Moreover, some data + specified in the input file of Vasp are transferred such as the mode + used to generate the **k** points, the line length (if line mode was + used), and any labels set for specific points. + + Examples + -------- + Read the **k** points data into a dictionary: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> calculation.kpoint.read() + {'mode': ..., 'line_length': ..., 'number_kpoints': ..., 'coordinates': array(...), 'weights': array(...)} + + Select the **k** points from the "kpoints_opt" mesh instead of the default one: + + >>> calculation.kpoint.read(selection="kpoints_opt") + {'mode': ..., 'line_length': ..., 'number_kpoints': ..., 'coordinates': array(...), 'weights': array(...), 'labels': ...} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.to_dict, + ) + + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + def line_length(self, selection=None) -> int: + """Get the number of points per line in the Brillouin zone. + + Returns + ------- + - + The number of points used to sample a single line. + + Examples + -------- + Get the number of points per line in the Brillouin zone: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> calculation.kpoint.line_length() + 48 + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.line_length, + ) + + def number_lines(self, selection=None) -> int: + """Get the number of lines in the Brillouin zone. + + Returns + ------- + - + The number of lines the band structure contains. For regular meshes this is + set to 1. + + Examples + -------- + Get the number of lines in the Brillouin zone for the "kpoints_opt" mesh: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> calculation.kpoint.number_lines(selection="kpoints_opt") + 4 + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.number_lines, + ) + + def number_kpoints(self, selection=None) -> int: + """Get the number of points in the Brillouin zone. + + Returns + ------- + - + The number of points used to sample the Brillouin zone. + + Examples + -------- + Get the number of points in the Brillouin zone: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> calculation.kpoint.number_kpoints() + 48 + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.number_kpoints, + ) + + def distances(self, selection=None) -> np.ndarray: + """Convert the coordinates of the **k** points into a one dimensional array. + + For every line in the Brillouin zone, the distance between each **k** point + and the start of the line is calculated. Then the distances of different + lines are concatenated into a single list. This routine is mostly useful + to plot data along high-symmetry lines like band structures. + + Returns + ------- + - + A reduction of the **k** points onto a one-dimensional array based + on the distance between the points. + + Examples + -------- + Convert the coordinates of the **k** points into a one dimensional array: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> calculation.kpoint.distances() + array([...]) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.distances, + ) + + def mode(self, selection=None) -> str: + """Get the **k**-point generation mode specified in the Vasp input file. + + Returns + ------- + - + A string representing which mode was used to setup the k-points. + + Examples + -------- + Get the **k**-point generation mode specified in the KPOINTS_OPT file: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> calculation.kpoint.mode(selection="kpoints_opt") + 'line' + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.mode, + ) + + def labels(self, selection=None) -> list[str] | None: + """Get any labels given in the input file for specific **k** points. + + The returned labels depend on the **k**-point mode and whether the user + provided explicit labels in the input file: + + - If labels are specified in the KPOINTS file, a list of strings is returned + with one entry per **k** point. Points with no user-defined label get an + empty string, while labeled points carry the name from the input file. + - If line mode is used but no labels were given, VASP automatically assigns + labels at the band edges. Interior points along the line receive an empty + string. Labels are formatted as LaTeX fractions. + - For any other mode without explicit labels (e.g., a regular Gamma or + Monkhorst-Pack grid), ``None`` is returned. + + Returns + ------- + - + A list of strings (one per **k** point) or ``None`` if no labeling is + applicable. + + Examples + -------- + If no labels were given and line mode is not used, returns None: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> result = calculation.kpoint.labels() + >>> assert result is None + + If line mode is used VASP automatically assigns labels to the band edges. In this case, + the method returns a list where band-edge points carry LaTeX-formatted coordinates and + interior points are empty strings. The example below uses the KPOINTS_OPT file: + + >>> calculation.kpoint.labels(selection="kpoints_opt") + ['$[0 0 0]$', ...] + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.labels, + ) + + def path_indices( + self, start: ArrayLike, finish: ArrayLike, selection=None + ) -> np.ndarray: + """Find linear dependent k points between start and finish. + + Loop over all possible k points and return the indices of the ones for which + k-point - start is linear dependent on finish - start. + + The primary use case is to extract a band-structure-like slice from a + regular **k**-point grid. In certain calculation types — such as + time-dependent DFT (TDDFT), GW, or BSE — VASP only supports uniform + Gamma or Monkhorst-Pack meshes. You can use this method to select all grid + points that happen to lie on a high-symmetry path. + + Parameters + ---------- + start + The starting **k** point of the path segment in fractional (crystal) + coordinates. Expects exactly 3 coordinates. + finish + The ending **k** point of the path segment in fractional (crystal) + coordinates. Expects exactly 3 coordinates. + + Returns + ------- + - + An integer array of indices (into the full **k**-point list) of all + **k** points that lie on the line segment from ``start`` to ``finish``. + + Examples + -------- + Extract all **k** points on a line through the Brillouin zone: + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + >>> start = [0, 0, 0.125] + >>> finish = [1, 0, 0.125] + >>> calculation.kpoint.path_indices(start, finish) + array([...]) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler.path_indices, + start, + finish, + ) + + def selections(self): + from py4vasp._raw import definition as raw_module + + return {self._quantity_name: list(raw_module.selections(self._quantity_name))} + + def _reciprocal_lattice_vectors(self, selection=None): + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + KpointHandler._reciprocal_lattice_vectors, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + KpointHandler.from_data, + KpointHandler.to_database, + ) + + +def _last_step(lattice_vectors): + if lattice_vectors.ndim == 2: + return lattice_vectors + else: + return lattice_vectors[-1] + + +def _line_distances(coordinates): + distances = np.zeros(len(coordinates)) + norms = np.linalg.norm(coordinates[1:] - coordinates[:-1], axis=1) + distances[1:] = np.cumsum(norms) + return distances + + +def _kpoint_label(kpoint): + fractions = [convert.Fraction(coordinate).latex() for coordinate in kpoint] + return f"$[{fractions[0]} {fractions[1]} {fractions[2]}]$" diff --git a/src/py4vasp/_calculation/local_moment.py b/src/py4vasp/_calculation/local_moment.py index 90820af74..307f38fed 100644 --- a/src/py4vasp/_calculation/local_moment.py +++ b/src/py4vasp/_calculation/local_moment.py @@ -1,785 +1,785 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy - -import numpy as np - -from py4vasp import _config, exception, raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, - slice_steps, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw.data_db import LocalMoment_DB -from py4vasp._third_party import view -from py4vasp._util import check, documentation, select - -_index_note = """\ -Notes ------ -The index order is different compared to the raw data when noncollinear calculations -are used. This routine returns the magnetic moments as (steps, orbitals, atoms, -directions).""" - -_moment_selection = """\ -selection : str - If VASP was run with LORBMOM = T, the orbital moments are computed and the routine - will default to the total moments. You can specify "spin" or "orbital" to select - the individual contributions instead. -""" - -_ORBITAL_PROJECTION = "orbital_projection" - - -class LocalMomentHandler: - """Handler for local moment data.""" - - length_moments = 1.5 - "Length in \u00c5 how a magnetic moment is displayed relative to the largest moment." - - def __init__(self, raw_local_moment: raw.LocalMoment, steps=None): - self._raw_local_moment = raw_local_moment - self._steps = steps - - @classmethod - def from_data( - cls, raw_local_moment: raw.LocalMoment, steps=None - ) -> "LocalMomentHandler": - return cls(raw_local_moment, steps=steps) - - def __str__(self) -> str: - if self._is_nonpolarized: - return "not spin polarized" - magmom = "MAGMOM = " - moments_last_step = self.magnetic("spin") - moments_to_string = lambda vec: " ".join(f"{moment:.2f}" for moment in vec) - if moments_last_step.ndim == 1: - return magmom + moments_to_string(moments_last_step) - else: - separator = " \\\n " - generator = (moments_to_string(vec) for vec in moments_last_step) - return magmom + separator.join(generator) - - def to_dict(self) -> dict: - return { - _ORBITAL_PROJECTION: self.selections()[_ORBITAL_PROJECTION], - "charge": self.projected_charge(), - **self._add_total_magnetic_moment(), - **self._add_spin_and_orbital_moments(), - } - - def to_database(self) -> dict: - spin_moments_orbitals = None - if not check.is_none(self._raw_local_moment.spin_moments): - if not self._is_nonpolarized: - spin_moments_orbitals = self._raw_local_moment.spin_moments[-1, -1] - spin_moment_total_min = None - spin_moment_total_max = None - if spin_moments_orbitals is not None: - spin_moments_total = np.sum(spin_moments_orbitals, axis=-1) - spin_moment_total_min = float(np.min(spin_moments_total)) - spin_moment_total_max = float(np.max(spin_moments_total)) - return LocalMoment_DB( - has_orbital_moments=self._has_orbital_moments, - final_spin_moment_total_min=spin_moment_total_min, - final_spin_moment_total_max=spin_moment_total_max, - ) - - def to_view(self, selection="total", supercell=None): - structure = StructureHandler.from_data( - self._raw_local_moment.structure, steps=self._steps - ) - viewer = structure.to_view(supercell) - if not self._is_nonpolarized: - viewer.ion_arrows = list( - self._prepare_magnetic_moments_for_plotting(selection) - ) - return viewer - - def projected_charge(self): - self._raise_error_if_steps_out_of_bounds() - return self._raw_local_moment.spin_moments[self._steps_or_last, 0] - - def projected_magnetic(self, selection="total"): - self._raise_error_if_steps_out_of_bounds() - self._raise_error_if_no_magnetic_moments() - tree = select.Tree.from_selection(selection) - moments = [self._magnetic_moments(sel) for sel in tree.selections()] - return np.squeeze(moments) - - def charge(self): - return _sum_over_orbitals(self.projected_charge()) - - def magnetic(self, selection="total"): - return _sum_over_orbitals( - self.projected_magnetic(selection), is_vector=self._is_noncollinear - ) - - def selections(self): - result = {} - if self._raw_local_moment.spin_moments.shape[-1] == 4: - result[_ORBITAL_PROJECTION] = ["s", "p", "d", "f"] - else: - result[_ORBITAL_PROJECTION] = ["s", "p", "d"] - if self._is_nonpolarized: - result["component"] = ["charge"] - elif self._has_orbital_moments: - result["component"] = ["charge", "total", "spin", "orbital"] - else: - result["component"] = ["charge", "total", "spin"] - return result - - def number_steps(self) -> int: - n = len(np.array(self._raw_local_moment.spin_moments)) - return len(range(n)[self._to_slice]) - - @property - def _is_nonpolarized(self): - return self._raw_local_moment.spin_moments.shape[1] == 1 - - @property - def _is_collinear(self): - return self._raw_local_moment.spin_moments.shape[1] == 2 - - @property - def _is_noncollinear(self): - return self._raw_local_moment.spin_moments.shape[1] == 4 - - @property - def _has_orbital_moments(self): - return not check.is_none(self._raw_local_moment.orbital_moments) - - @property - def _steps_or_last(self): - if self._steps is None or self._steps == -1: - return -1 - return self._steps - - @property - def _to_slice(self): - if self._steps is None or self._steps == -1: - return slice(-1, None) - if isinstance(self._steps, slice): - return self._steps - return slice(self._steps, self._steps + 1) - - def _magnetic_moments(self, selection): - self._raise_error_if_selection_not_available(selection) - if self._is_collinear: - return self._spin_moments() - else: - return self._noncollinear_moments(selection[0]) - - def _noncollinear_moments(self, selection): - spin_moments = self._spin_moments() - orbital_moments = self._orbital_moments(spin_moments) - if selection == "orbital": - moments = orbital_moments - elif selection == "spin": - moments = spin_moments - else: - moments = spin_moments + orbital_moments - direction_axis = 1 if moments.ndim == 4 else 0 - return np.moveaxis(moments, direction_axis, -1) - - def _spin_moments(self): - return self._raw_local_moment.spin_moments[self._steps_or_last, 1:] - - def _orbital_moments(self, spin_moments): - if not self._has_orbital_moments: - return np.zeros_like(spin_moments) - zero_s_moments = np.zeros((*spin_moments.shape[:-1], 1)) - orbital_moments = self._raw_local_moment.orbital_moments[self._steps_or_last] - return np.concatenate((zero_s_moments, orbital_moments), axis=-1) - - def _add_total_magnetic_moment(self): - if self._is_nonpolarized: - return {} - return {"total": self.projected_magnetic()} - - def _add_spin_and_orbital_moments(self): - if not self._has_orbital_moments: - return {} - spin_moments = self._spin_moments() - orbital_moments = self._orbital_moments(spin_moments) - direction_axis = 1 if spin_moments.ndim == 4 else 0 - return { - "spin": np.moveaxis(spin_moments, direction_axis, -1), - "orbital": np.moveaxis(orbital_moments, direction_axis, -1), - } - - def _prepare_magnetic_moments_for_plotting(self, selection): - tree = select.Tree.from_selection(selection) - for (sel, *_) in tree.selections(): - moments = self.magnetic(sel) - moments = self._make_sure_moments_have_timestep_dimension(moments) - moments = _convert_moment_to_3d_vector(moments) - max_length_moments = _max_length_moments(moments) - if max_length_moments > 1e-15: - rescale_moments = LocalMomentHandler.length_moments / max_length_moments - yield view.IonArrow( - quantity=rescale_moments * moments, - label=f"{sel} moments", - color=_color(sel), - radius=0.2, - ) - - def _make_sure_moments_have_timestep_dimension(self, moments): - is_slice = isinstance(self._steps, slice) - if not is_slice and moments is not None: - moments = moments[np.newaxis] - return moments - - def _raise_error_if_steps_out_of_bounds(self): - try: - np.zeros(self._raw_local_moment.spin_moments.shape[0])[self._steps_or_last] - except IndexError as error: - raise exception.IncorrectUsage( - f"Error reading the magnetic moments. Please check if the steps " - f"`{self._steps}` are properly formatted and within the boundaries." - ) from error - - def _raise_error_if_no_magnetic_moments(self): - if self._is_nonpolarized: - raise exception.NoData( - "There are no magnetic moments in the data. Please make sure that you " - "either set ISPIN = 2 or LNONCOLLINEAR = T or LSORBIT = T." - ) - - def _raise_error_if_selection_not_available(self, selection): - if len(selection) != 1: - raise exception.IncorrectUsage() - selection = selection[0] - if selection not in ("spin", "orbital", "total"): - raise exception.IncorrectUsage( - f"The selection {selection} is incorrect. Please check if it is spelled " - "correctly. Possible choices are total, spin, or orbital." - ) - if selection != "orbital" or self._has_orbital_moments: - return - raise exception.NoData( - "There are no orbital moments in the VASP output. Please make sure that you " - "run the calculation with LORBMOM = T and LSORBIT = T." - ) - - -@quantity("local_moment") -class LocalMoment(view.Mixin): - """The local moments describe the charge and magnetization near an atom. - - The projection on local moments is particularly relevant in the context of - magnetic materials. It analyzes the electronic states in the vicinity of an - atom by projecting the electronic orbitals onto the localized projectors of - the PAWs. The local moments help understanding the magnetic ordering, the spin - polarization, and the influence of neighboring atoms on the magnetic behavior. - - This class allows to access the computed moments from a VASP calculation. - Remember that VASP calculates the projections only if you need to set - :tag:`LORBIT` in the INCAR file. If the system is computed without spin - polarization, the resulting moments correspond only to the local charges - resolved by angular momentum. For collinear calculation, additionally the - magnetic moment are computed. In noncollinear calculations, the magnetization - becomes a vector. When comparing the results extracted from VASP to experimental - observation, please be aware that the finite size of the radius in the projection - may influence the observed moments. Hence, there is no one-to-one correspondence - to the experimental moments. - - Examples - -------- - Let us create some example data so that we can illustrate how to use this class. - Of course you can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path, "collinear") - - If you access the local moments, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.local_moment.number_steps() - 1 - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.local_moment[:].number_steps() - 4 - - You can also select specific steps or a subset of steps as follows - - >>> calculation.local_moment[3].number_steps() - 1 - >>> calculation.local_moment[1:4].number_steps() - 3 - """ - - length_moments = LocalMomentHandler.length_moments - - def __init__(self, source, quantity_name: str = "local_moment", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_local_moment: raw.LocalMoment) -> "LocalMoment": - return cls(source=DataSource(raw_local_moment)) - - def __getitem__(self, steps) -> "LocalMoment": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw): - return LocalMomentHandler.from_data(raw, steps=self._steps) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - LocalMomentHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - @documentation.format(index_note=_index_note) - def read(self) -> dict: - """Read the charges and magnetization data into a dictionary. - - Be careful when comparing the magnetic moments to experimental data. The - finite size of the projection sphere may influence the observed moments. Hence, - there is no one-to-one correspondence to the experimental moments. - - Returns - ------- - dict - Contains the charges and magnetic moments generated by VASP projected - on atoms and orbitals. - - {index_note} - - Examples - -------- - First we create some example data so that we can illustrate how to use this - method. Of course you can also use your own VASP calculation data if you have - it available. - - >>> from py4vasp import demo - >>> collinear_calculation = demo.calculation(path, "collinear") - >>> noncollinear_calculation = demo.calculation(path, "noncollinear") - - If you use the `read` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> collinear_calculation.local_moment.read() - {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[...]]), - 'total': array([[...]])}} - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the charge and the total moment contain an additional - dimension for the different steps. - - >>> collinear_calculation.local_moment[:].read() - {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[[...]]]), - 'total': array([[[...]]])}} - - You can also select specific steps or a subset of steps as follows - - >>> collinear_calculation.local_moment[2].read() - {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[...]]), - 'total': array([[...]])}} - >>> collinear_calculation.local_moment[0:3].read() - {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[[...]]]), - 'total': array([[[...]]])}} - - For noncollinear calculations, the magnetic moments are vectors. In addition, - if the calculation was run with LORBMOM = T, orbital and spin moments are - reported separately. - - >>> noncollinear_calculation.local_moment.read() - {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[...]]), - 'total': array([[[...]]]), 'spin': array([[[...]]]), 'orbital': array([[[...]]])}} - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - LocalMomentHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - @documentation.format(selection=_moment_selection) - def to_view(self, selection="total", supercell=None): - """Visualize the magnetic moments as arrows inside the structure. - - Be aware that the magnetic moments are rescaled for better visibility. The length - of the largest moment is set to :attr:`length_moments`. If your moments are - vanishingly small, this may lead to unexpectedly large arrows. Please double check - the actual values with :meth:`magnetic`. - - Parameters - ---------- - {selection} - - Returns - ------- - View - Contains the atoms and the unit cell as well as an arrow indicating the - strength of the magnetic moment. If noncollinear magnetism is used the - moment points in the actual direction; for collinear magnetism the - moments are aligned along the z axis by convention. - - Examples - -------- - First we create some example data so that we can illustrate how to use this - method. Of course you can also use your own VASP calculation data if you have - it available. - - >>> from py4vasp import demo - >>> collinear_calculation = demo.calculation(path, "collinear") - >>> noncollinear_calculation = demo.calculation(path, "noncollinear") - - If you use the `to_view` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> collinear_calculation.local_moment.to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) - - For collinear calculations, the magnetic moments are scalars aligned along the - z axis. You can see this from the arrows pointing either up or down or x and y - components being zero. - - To select the results for all steps, you don't specify the array boundaries. - - >>> collinear_calculation.local_moment[:].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) - - You can also select specific steps or a subset of steps as follows - - >>> collinear_calculation.local_moment[2].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) - >>> collinear_calculation.local_moment[0:3].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) - - For noncollinear calculations, the magnetic moments are aligned according to - the spin axis (:tag:`SAXIS`). The view has the same interface, but the arrows now - point in the actual direction of the magnetic moments. - - >>> noncollinear_calculation.local_moment.to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) - - You can also select spin or orbital moments separately. - - >>> noncollinear_calculation.local_moment.to_view(selection="spin, orbital") - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='spin moments', ...), IonArrow(quantity=array([[[...]]]), label='orbital moments', ...)], ...) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - LocalMomentHandler.to_view, - supercell, - ) - - def projected_charge(self): - """Read the orbital- and site-projected charges of the selected steps. - - Returns - ------- - np.ndarray - Contains the charges for the selected steps projected on atoms and orbitals. - - Examples - -------- - First we create some example data so that we can illustrate how to use this - method. Of course you can also use your own VASP calculation data if you have - it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path, "collinear") - - If you use the `projected_charge` method, the result will depend on the steps - that you selected with the [] operator. Without any selection the results from - the final step will be used. - - >>> calculation.local_moment.projected_charge() - array([[...]]) - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the charge contains an additional dimension for the - different steps. - - >>> calculation.local_moment[:].projected_charge() - array([[[...]]]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - LocalMomentHandler.projected_charge, - ) - - @documentation.format(selection=_moment_selection, index_note=_index_note) - def projected_magnetic(self, selection="total"): - """Read the orbital- and site-projected magnetic moments of the selected steps. - - Parameters - ---------- - {selection} - - Returns - ------- - np.ndarray - Contains the magnetic moments for the selected steps projected on atoms - and orbitals. - - {index_note} - - Examples - -------- - First we create some example data so that we can illustrate how to use this - method. Of course you can also use your own VASP calculation data if you have - it available. - - >>> from py4vasp import demo - >>> collinear_calculation = demo.calculation(path, "collinear") - >>> noncollinear_calculation = demo.calculation(path, "noncollinear") - - If you use the `projected_magnetic` method, the result will depend on the steps - that you selected with the [] operator. Without any selection the results from - the final step will be used. - - >>> collinear_calculation.local_moment.projected_magnetic() - array([[...]]) - - To select the results for all steps, you don't specify the array boundaries. - - >>> collinear_calculation.local_moment[:].projected_magnetic() - array([[[...]]]) - - You can also select specific steps or a subset of steps as follows - - >>> collinear_calculation.local_moment[2].projected_magnetic() - array([[...]]) - >>> collinear_calculation.local_moment[0:3].projected_magnetic() - array([[[...]]]) - - For noncollinear calculations, the magnetic moments are aligned according to - the spin axis (:tag:`SAXIS`). Since the moments are now vectors, the result - has an additional dimension for the different directions. - - >>> noncollinear_calculation.local_moment.projected_magnetic() - array([[[...]]]) - - You can also select spin or orbital moments separately. - - >>> noncollinear_calculation.local_moment.projected_magnetic(selection="spin") - array([[[...]]]) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - LocalMomentHandler.projected_magnetic, - ) - - def charge(self): - """Read the site-projected charges of the selected steps. - - Returns - ------- - np.ndarray - Contains the charges for the selected steps projected on atoms. Equivalent - to summing the projected charges over the orbitals. - - Examples - -------- - First we create some example data so that we can illustrate how to use this - method. Of course you can also use your own VASP calculation data if you have - it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path, "collinear") - - If you use the `charge` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.local_moment.charge() - array([...]) - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the charge contains an additional dimension for the - different steps. - - >>> calculation.local_moment[:].charge() - array([[...]]) - - The charge is equivalent to summing the projected charges over the orbitals - - >>> np.allclose(calculation.local_moment.charge(), - ... calculation.local_moment.projected_charge().sum(axis=-1)) - True - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - LocalMomentHandler.charge, - ) - - @documentation.format(selection=_moment_selection, index_note=_index_note) - def magnetic(self, selection="total"): - """Read the site-projected magnetic moments of the selected steps. - - Be careful when comparing the magnetic moments to experimental data. The - finite size of the projection sphere may influence the observed moments. Hence, - there is no one-to-one correspondence to the experimental moments. - - Parameters - ---------- - {selection} - - Returns - ------- - np.ndarray - Contains the magnetic moments for the selected steps projected on atoms. - Equivalent to summing the projected magnetic moments over the orbitals. - - {index_note} - - Examples - -------- - First we create some example data so that we can illustrate how to use this - method. Of course you can also use your own VASP calculation data if you have - it available. - - >>> from py4vasp import demo - >>> collinear_calculation = demo.calculation(path, "collinear") - >>> noncollinear_calculation = demo.calculation(path, "noncollinear") - - If you use the `magnetic` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> collinear_calculation.local_moment.magnetic() - array([...]) - - To select the results for all steps, you don't specify the array boundaries. - - >>> collinear_calculation.local_moment[:].magnetic() - array([[...]]) - - You can also select specific steps or a subset of steps as follows - - >>> collinear_calculation.local_moment[2].magnetic() - array([...]) - >>> collinear_calculation.local_moment[0:3].magnetic() - array([[...]]) - - For noncollinear calculations, the magnetic moments are aligned according to - the spin axis (:tag:`SAXIS`). Since the moments are now vectors, the result - has an additional dimension for the different directions. - - >>> noncollinear_calculation.local_moment.magnetic() - array([[...]]) - - You can also select spin or orbital moments separately. - - >>> noncollinear_calculation.local_moment.magnetic(selection="spin") - array([[...]]) - - The magnetic moment is equivalent to summing the projected magnetic moments - over the orbitals - - >>> np.allclose(collinear_calculation.local_moment.magnetic(), - ... collinear_calculation.local_moment.projected_magnetic().sum(axis=-1)) - True - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - LocalMomentHandler.magnetic, - ) - - def selections(self) -> dict: - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - LocalMomentHandler.selections, - ) - - def number_steps(self) -> int: - """Return the number of local moments in the trajectory.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - LocalMomentHandler.number_steps, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - LocalMomentHandler.from_data, - LocalMomentHandler.to_database, - ) - - -def _sum_over_orbitals(quantity, is_vector=False): - if quantity is None: - return None - if is_vector: - return np.sum(quantity, axis=-2) - return np.sum(quantity, axis=-1) - - -def _convert_moment_to_3d_vector(moments): - if moments is not None and moments.ndim == 2: - moments = moments.reshape((*moments.shape, 1)) - no_new_moments = (0, 0) - add_zero_for_xy_axis = (2, 0) - padding = (no_new_moments, no_new_moments, add_zero_for_xy_axis) - moments = np.pad(moments, padding) - return moments - - -def _max_length_moments(moments): - if moments is not None: - return np.max(np.linalg.norm(moments, axis=2)) - else: - return 0.0 - - -def _color(selection): - if selection == "total": - return _config.VASP_COLORS["blue"] - if selection == "spin": - return _config.VASP_COLORS["purple"] - if selection == "orbital": - return _config.VASP_COLORS["red"] - raise exception.IncorrectUsage(f"Unknown component {selection} selected.") +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy + +import numpy as np + +from py4vasp import _config, exception, raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, + slice_steps, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw.data_db import LocalMoment_DB +from py4vasp._third_party import view +from py4vasp._util import check, documentation, select + +_index_note = """\ +Notes +----- +The index order is different compared to the raw data when noncollinear calculations +are used. This routine returns the magnetic moments as (steps, orbitals, atoms, +directions).""" + +_moment_selection = """\ +selection : str + If VASP was run with LORBMOM = T, the orbital moments are computed and the routine + will default to the total moments. You can specify "spin" or "orbital" to select + the individual contributions instead. +""" + +_ORBITAL_PROJECTION = "orbital_projection" + + +class LocalMomentHandler: + """Handler for local moment data.""" + + length_moments = 1.5 + "Length in \u00c5 how a magnetic moment is displayed relative to the largest moment." + + def __init__(self, raw_local_moment: raw.LocalMoment, steps=None): + self._raw_local_moment = raw_local_moment + self._steps = steps + + @classmethod + def from_data( + cls, raw_local_moment: raw.LocalMoment, steps=None + ) -> "LocalMomentHandler": + return cls(raw_local_moment, steps=steps) + + def __str__(self) -> str: + if self._is_nonpolarized: + return "not spin polarized" + magmom = "MAGMOM = " + moments_last_step = self.magnetic("spin") + moments_to_string = lambda vec: " ".join(f"{moment:.2f}" for moment in vec) + if moments_last_step.ndim == 1: + return magmom + moments_to_string(moments_last_step) + else: + separator = " \\\n " + generator = (moments_to_string(vec) for vec in moments_last_step) + return magmom + separator.join(generator) + + def to_dict(self) -> dict: + return { + _ORBITAL_PROJECTION: self.selections()[_ORBITAL_PROJECTION], + "charge": self.projected_charge(), + **self._add_total_magnetic_moment(), + **self._add_spin_and_orbital_moments(), + } + + def to_database(self) -> dict: + spin_moments_orbitals = None + if not check.is_none(self._raw_local_moment.spin_moments): + if not self._is_nonpolarized: + spin_moments_orbitals = self._raw_local_moment.spin_moments[-1, -1] + spin_moment_total_min = None + spin_moment_total_max = None + if spin_moments_orbitals is not None: + spin_moments_total = np.sum(spin_moments_orbitals, axis=-1) + spin_moment_total_min = float(np.min(spin_moments_total)) + spin_moment_total_max = float(np.max(spin_moments_total)) + return LocalMoment_DB( + has_orbital_moments=self._has_orbital_moments, + final_spin_moment_total_min=spin_moment_total_min, + final_spin_moment_total_max=spin_moment_total_max, + ) + + def to_view(self, selection="total", supercell=None): + structure = StructureHandler.from_data( + self._raw_local_moment.structure, steps=self._steps + ) + viewer = structure.to_view(supercell) + if not self._is_nonpolarized: + viewer.ion_arrows = list( + self._prepare_magnetic_moments_for_plotting(selection) + ) + return viewer + + def projected_charge(self): + self._raise_error_if_steps_out_of_bounds() + return self._raw_local_moment.spin_moments[self._steps_or_last, 0] + + def projected_magnetic(self, selection="total"): + self._raise_error_if_steps_out_of_bounds() + self._raise_error_if_no_magnetic_moments() + tree = select.Tree.from_selection(selection) + moments = [self._magnetic_moments(sel) for sel in tree.selections()] + return np.squeeze(moments) + + def charge(self): + return _sum_over_orbitals(self.projected_charge()) + + def magnetic(self, selection="total"): + return _sum_over_orbitals( + self.projected_magnetic(selection), is_vector=self._is_noncollinear + ) + + def selections(self): + result = {} + if self._raw_local_moment.spin_moments.shape[-1] == 4: + result[_ORBITAL_PROJECTION] = ["s", "p", "d", "f"] + else: + result[_ORBITAL_PROJECTION] = ["s", "p", "d"] + if self._is_nonpolarized: + result["component"] = ["charge"] + elif self._has_orbital_moments: + result["component"] = ["charge", "total", "spin", "orbital"] + else: + result["component"] = ["charge", "total", "spin"] + return result + + def number_steps(self) -> int: + n = len(np.array(self._raw_local_moment.spin_moments)) + return len(range(n)[self._to_slice]) + + @property + def _is_nonpolarized(self): + return self._raw_local_moment.spin_moments.shape[1] == 1 + + @property + def _is_collinear(self): + return self._raw_local_moment.spin_moments.shape[1] == 2 + + @property + def _is_noncollinear(self): + return self._raw_local_moment.spin_moments.shape[1] == 4 + + @property + def _has_orbital_moments(self): + return not check.is_none(self._raw_local_moment.orbital_moments) + + @property + def _steps_or_last(self): + if self._steps is None or self._steps == -1: + return -1 + return self._steps + + @property + def _to_slice(self): + if self._steps is None or self._steps == -1: + return slice(-1, None) + if isinstance(self._steps, slice): + return self._steps + return slice(self._steps, self._steps + 1) + + def _magnetic_moments(self, selection): + self._raise_error_if_selection_not_available(selection) + if self._is_collinear: + return self._spin_moments() + else: + return self._noncollinear_moments(selection[0]) + + def _noncollinear_moments(self, selection): + spin_moments = self._spin_moments() + orbital_moments = self._orbital_moments(spin_moments) + if selection == "orbital": + moments = orbital_moments + elif selection == "spin": + moments = spin_moments + else: + moments = spin_moments + orbital_moments + direction_axis = 1 if moments.ndim == 4 else 0 + return np.moveaxis(moments, direction_axis, -1) + + def _spin_moments(self): + return self._raw_local_moment.spin_moments[self._steps_or_last, 1:] + + def _orbital_moments(self, spin_moments): + if not self._has_orbital_moments: + return np.zeros_like(spin_moments) + zero_s_moments = np.zeros((*spin_moments.shape[:-1], 1)) + orbital_moments = self._raw_local_moment.orbital_moments[self._steps_or_last] + return np.concatenate((zero_s_moments, orbital_moments), axis=-1) + + def _add_total_magnetic_moment(self): + if self._is_nonpolarized: + return {} + return {"total": self.projected_magnetic()} + + def _add_spin_and_orbital_moments(self): + if not self._has_orbital_moments: + return {} + spin_moments = self._spin_moments() + orbital_moments = self._orbital_moments(spin_moments) + direction_axis = 1 if spin_moments.ndim == 4 else 0 + return { + "spin": np.moveaxis(spin_moments, direction_axis, -1), + "orbital": np.moveaxis(orbital_moments, direction_axis, -1), + } + + def _prepare_magnetic_moments_for_plotting(self, selection): + tree = select.Tree.from_selection(selection) + for (sel, *_) in tree.selections(): + moments = self.magnetic(sel) + moments = self._make_sure_moments_have_timestep_dimension(moments) + moments = _convert_moment_to_3d_vector(moments) + max_length_moments = _max_length_moments(moments) + if max_length_moments > 1e-15: + rescale_moments = LocalMomentHandler.length_moments / max_length_moments + yield view.IonArrow( + quantity=rescale_moments * moments, + label=f"{sel} moments", + color=_color(sel), + radius=0.2, + ) + + def _make_sure_moments_have_timestep_dimension(self, moments): + is_slice = isinstance(self._steps, slice) + if not is_slice and moments is not None: + moments = moments[np.newaxis] + return moments + + def _raise_error_if_steps_out_of_bounds(self): + try: + np.zeros(self._raw_local_moment.spin_moments.shape[0])[self._steps_or_last] + except IndexError as error: + raise exception.IncorrectUsage( + f"Error reading the magnetic moments. Please check if the steps " + f"`{self._steps}` are properly formatted and within the boundaries." + ) from error + + def _raise_error_if_no_magnetic_moments(self): + if self._is_nonpolarized: + raise exception.NoData( + "There are no magnetic moments in the data. Please make sure that you " + "either set ISPIN = 2 or LNONCOLLINEAR = T or LSORBIT = T." + ) + + def _raise_error_if_selection_not_available(self, selection): + if len(selection) != 1: + raise exception.IncorrectUsage() + selection = selection[0] + if selection not in ("spin", "orbital", "total"): + raise exception.IncorrectUsage( + f"The selection {selection} is incorrect. Please check if it is spelled " + "correctly. Possible choices are total, spin, or orbital." + ) + if selection != "orbital" or self._has_orbital_moments: + return + raise exception.NoData( + "There are no orbital moments in the VASP output. Please make sure that you " + "run the calculation with LORBMOM = T and LSORBIT = T." + ) + + +@quantity("local_moment") +class LocalMoment(view.Mixin): + """The local moments describe the charge and magnetization near an atom. + + The projection on local moments is particularly relevant in the context of + magnetic materials. It analyzes the electronic states in the vicinity of an + atom by projecting the electronic orbitals onto the localized projectors of + the PAWs. The local moments help understanding the magnetic ordering, the spin + polarization, and the influence of neighboring atoms on the magnetic behavior. + + This class allows to access the computed moments from a VASP calculation. + Remember that VASP calculates the projections only if you need to set + :tag:`LORBIT` in the INCAR file. If the system is computed without spin + polarization, the resulting moments correspond only to the local charges + resolved by angular momentum. For collinear calculation, additionally the + magnetic moment are computed. In noncollinear calculations, the magnetization + becomes a vector. When comparing the results extracted from VASP to experimental + observation, please be aware that the finite size of the radius in the projection + may influence the observed moments. Hence, there is no one-to-one correspondence + to the experimental moments. + + Examples + -------- + Let us create some example data so that we can illustrate how to use this class. + Of course you can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path, "collinear") + + If you access the local moments, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.local_moment.number_steps() + 1 + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.local_moment[:].number_steps() + 4 + + You can also select specific steps or a subset of steps as follows + + >>> calculation.local_moment[3].number_steps() + 1 + >>> calculation.local_moment[1:4].number_steps() + 3 + """ + + length_moments = LocalMomentHandler.length_moments + + def __init__(self, source, quantity_name: str = "local_moment", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_local_moment: raw.LocalMoment) -> "LocalMoment": + return cls(source=DataSource(raw_local_moment)) + + def __getitem__(self, steps) -> "LocalMoment": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw): + return LocalMomentHandler.from_data(raw, steps=self._steps) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + LocalMomentHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + @documentation.format(index_note=_index_note) + def read(self) -> dict: + """Read the charges and magnetization data into a dictionary. + + Be careful when comparing the magnetic moments to experimental data. The + finite size of the projection sphere may influence the observed moments. Hence, + there is no one-to-one correspondence to the experimental moments. + + Returns + ------- + dict + Contains the charges and magnetic moments generated by VASP projected + on atoms and orbitals. + + {index_note} + + Examples + -------- + First we create some example data so that we can illustrate how to use this + method. Of course you can also use your own VASP calculation data if you have + it available. + + >>> from py4vasp import demo + >>> collinear_calculation = demo.calculation(path, "collinear") + >>> noncollinear_calculation = demo.calculation(path, "noncollinear") + + If you use the `read` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> collinear_calculation.local_moment.read() + {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[...]]), + 'total': array([[...]])}} + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the charge and the total moment contain an additional + dimension for the different steps. + + >>> collinear_calculation.local_moment[:].read() + {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[[...]]]), + 'total': array([[[...]]])}} + + You can also select specific steps or a subset of steps as follows + + >>> collinear_calculation.local_moment[2].read() + {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[...]]), + 'total': array([[...]])}} + >>> collinear_calculation.local_moment[0:3].read() + {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[[...]]]), + 'total': array([[[...]]])}} + + For noncollinear calculations, the magnetic moments are vectors. In addition, + if the calculation was run with LORBMOM = T, orbital and spin moments are + reported separately. + + >>> noncollinear_calculation.local_moment.read() + {{'orbital_projection': ['s', 'p', 'd'], 'charge': array([[...]]), + 'total': array([[[...]]]), 'spin': array([[[...]]]), 'orbital': array([[[...]]])}} + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + LocalMomentHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + @documentation.format(selection=_moment_selection) + def to_view(self, selection="total", supercell=None): + """Visualize the magnetic moments as arrows inside the structure. + + Be aware that the magnetic moments are rescaled for better visibility. The length + of the largest moment is set to :attr:`length_moments`. If your moments are + vanishingly small, this may lead to unexpectedly large arrows. Please double check + the actual values with :meth:`magnetic`. + + Parameters + ---------- + {selection} + + Returns + ------- + View + Contains the atoms and the unit cell as well as an arrow indicating the + strength of the magnetic moment. If noncollinear magnetism is used the + moment points in the actual direction; for collinear magnetism the + moments are aligned along the z axis by convention. + + Examples + -------- + First we create some example data so that we can illustrate how to use this + method. Of course you can also use your own VASP calculation data if you have + it available. + + >>> from py4vasp import demo + >>> collinear_calculation = demo.calculation(path, "collinear") + >>> noncollinear_calculation = demo.calculation(path, "noncollinear") + + If you use the `to_view` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> collinear_calculation.local_moment.to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) + + For collinear calculations, the magnetic moments are scalars aligned along the + z axis. You can see this from the arrows pointing either up or down or x and y + components being zero. + + To select the results for all steps, you don't specify the array boundaries. + + >>> collinear_calculation.local_moment[:].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) + + You can also select specific steps or a subset of steps as follows + + >>> collinear_calculation.local_moment[2].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) + >>> collinear_calculation.local_moment[0:3].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) + + For noncollinear calculations, the magnetic moments are aligned according to + the spin axis (:tag:`SAXIS`). The view has the same interface, but the arrows now + point in the actual direction of the magnetic moments. + + >>> noncollinear_calculation.local_moment.to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='total moments', ...)], ...) + + You can also select spin or orbital moments separately. + + >>> noncollinear_calculation.local_moment.to_view(selection="spin, orbital") + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='spin moments', ...), IonArrow(quantity=array([[[...]]]), label='orbital moments', ...)], ...) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + LocalMomentHandler.to_view, + supercell, + ) + + def projected_charge(self): + """Read the orbital- and site-projected charges of the selected steps. + + Returns + ------- + np.ndarray + Contains the charges for the selected steps projected on atoms and orbitals. + + Examples + -------- + First we create some example data so that we can illustrate how to use this + method. Of course you can also use your own VASP calculation data if you have + it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path, "collinear") + + If you use the `projected_charge` method, the result will depend on the steps + that you selected with the [] operator. Without any selection the results from + the final step will be used. + + >>> calculation.local_moment.projected_charge() + array([[...]]) + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the charge contains an additional dimension for the + different steps. + + >>> calculation.local_moment[:].projected_charge() + array([[[...]]]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + LocalMomentHandler.projected_charge, + ) + + @documentation.format(selection=_moment_selection, index_note=_index_note) + def projected_magnetic(self, selection="total"): + """Read the orbital- and site-projected magnetic moments of the selected steps. + + Parameters + ---------- + {selection} + + Returns + ------- + np.ndarray + Contains the magnetic moments for the selected steps projected on atoms + and orbitals. + + {index_note} + + Examples + -------- + First we create some example data so that we can illustrate how to use this + method. Of course you can also use your own VASP calculation data if you have + it available. + + >>> from py4vasp import demo + >>> collinear_calculation = demo.calculation(path, "collinear") + >>> noncollinear_calculation = demo.calculation(path, "noncollinear") + + If you use the `projected_magnetic` method, the result will depend on the steps + that you selected with the [] operator. Without any selection the results from + the final step will be used. + + >>> collinear_calculation.local_moment.projected_magnetic() + array([[...]]) + + To select the results for all steps, you don't specify the array boundaries. + + >>> collinear_calculation.local_moment[:].projected_magnetic() + array([[[...]]]) + + You can also select specific steps or a subset of steps as follows + + >>> collinear_calculation.local_moment[2].projected_magnetic() + array([[...]]) + >>> collinear_calculation.local_moment[0:3].projected_magnetic() + array([[[...]]]) + + For noncollinear calculations, the magnetic moments are aligned according to + the spin axis (:tag:`SAXIS`). Since the moments are now vectors, the result + has an additional dimension for the different directions. + + >>> noncollinear_calculation.local_moment.projected_magnetic() + array([[[...]]]) + + You can also select spin or orbital moments separately. + + >>> noncollinear_calculation.local_moment.projected_magnetic(selection="spin") + array([[[...]]]) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + LocalMomentHandler.projected_magnetic, + ) + + def charge(self): + """Read the site-projected charges of the selected steps. + + Returns + ------- + np.ndarray + Contains the charges for the selected steps projected on atoms. Equivalent + to summing the projected charges over the orbitals. + + Examples + -------- + First we create some example data so that we can illustrate how to use this + method. Of course you can also use your own VASP calculation data if you have + it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path, "collinear") + + If you use the `charge` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.local_moment.charge() + array([...]) + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the charge contains an additional dimension for the + different steps. + + >>> calculation.local_moment[:].charge() + array([[...]]) + + The charge is equivalent to summing the projected charges over the orbitals + + >>> np.allclose(calculation.local_moment.charge(), + ... calculation.local_moment.projected_charge().sum(axis=-1)) + True + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + LocalMomentHandler.charge, + ) + + @documentation.format(selection=_moment_selection, index_note=_index_note) + def magnetic(self, selection="total"): + """Read the site-projected magnetic moments of the selected steps. + + Be careful when comparing the magnetic moments to experimental data. The + finite size of the projection sphere may influence the observed moments. Hence, + there is no one-to-one correspondence to the experimental moments. + + Parameters + ---------- + {selection} + + Returns + ------- + np.ndarray + Contains the magnetic moments for the selected steps projected on atoms. + Equivalent to summing the projected magnetic moments over the orbitals. + + {index_note} + + Examples + -------- + First we create some example data so that we can illustrate how to use this + method. Of course you can also use your own VASP calculation data if you have + it available. + + >>> from py4vasp import demo + >>> collinear_calculation = demo.calculation(path, "collinear") + >>> noncollinear_calculation = demo.calculation(path, "noncollinear") + + If you use the `magnetic` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> collinear_calculation.local_moment.magnetic() + array([...]) + + To select the results for all steps, you don't specify the array boundaries. + + >>> collinear_calculation.local_moment[:].magnetic() + array([[...]]) + + You can also select specific steps or a subset of steps as follows + + >>> collinear_calculation.local_moment[2].magnetic() + array([...]) + >>> collinear_calculation.local_moment[0:3].magnetic() + array([[...]]) + + For noncollinear calculations, the magnetic moments are aligned according to + the spin axis (:tag:`SAXIS`). Since the moments are now vectors, the result + has an additional dimension for the different directions. + + >>> noncollinear_calculation.local_moment.magnetic() + array([[...]]) + + You can also select spin or orbital moments separately. + + >>> noncollinear_calculation.local_moment.magnetic(selection="spin") + array([[...]]) + + The magnetic moment is equivalent to summing the projected magnetic moments + over the orbitals + + >>> np.allclose(collinear_calculation.local_moment.magnetic(), + ... collinear_calculation.local_moment.projected_magnetic().sum(axis=-1)) + True + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + LocalMomentHandler.magnetic, + ) + + def selections(self) -> dict: + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + LocalMomentHandler.selections, + ) + + def number_steps(self) -> int: + """Return the number of local moments in the trajectory.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + LocalMomentHandler.number_steps, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + LocalMomentHandler.from_data, + LocalMomentHandler.to_database, + ) + + +def _sum_over_orbitals(quantity, is_vector=False): + if quantity is None: + return None + if is_vector: + return np.sum(quantity, axis=-2) + return np.sum(quantity, axis=-1) + + +def _convert_moment_to_3d_vector(moments): + if moments is not None and moments.ndim == 2: + moments = moments.reshape((*moments.shape, 1)) + no_new_moments = (0, 0) + add_zero_for_xy_axis = (2, 0) + padding = (no_new_moments, no_new_moments, add_zero_for_xy_axis) + moments = np.pad(moments, padding) + return moments + + +def _max_length_moments(moments): + if moments is not None: + return np.max(np.linalg.norm(moments, axis=2)) + else: + return 0.0 + + +def _color(selection): + if selection == "total": + return _config.VASP_COLORS["blue"] + if selection == "spin": + return _config.VASP_COLORS["purple"] + if selection == "orbital": + return _config.VASP_COLORS["red"] + raise exception.IncorrectUsage(f"Unknown component {selection} selected.") diff --git a/src/py4vasp/_calculation/nics.py b/src/py4vasp/_calculation/nics.py index a6272613d..7a5cda642 100644 --- a/src/py4vasp/_calculation/nics.py +++ b/src/py4vasp/_calculation/nics.py @@ -1,502 +1,502 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy -from typing import Optional, Union - -import numpy as np - -from py4vasp import _config, exception -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Nics_DB -from py4vasp._third_party import graph, view -from py4vasp._util import check, documentation, import_, index, select, slicing - -pretty = import_.optional("IPython.lib.pretty") - -_DEFAULT_SELECTION: str = "isotropic" - - -class NicsHandler: - """Handler for NICS data — performs all data access and transformation.""" - - def __init__(self, raw_nics: raw.Nics): - self._raw_nics = raw_nics - - @classmethod - def from_data(cls, raw_nics: raw.Nics) -> "NicsHandler": - return cls(raw_nics) - - def __str__(self) -> str: - raw_stoichiometry = self._raw_nics.structure.stoichiometry - stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry) - if self._data_is_on_grid: - data_string = self._grid_to_string() - else: - data_string = self._points_to_string() - return f"""\ -nucleus-independent chemical shift: - structure: {pretty.pretty(stoichiometry)} -{data_string}""" - - def to_dict(self) -> dict: - """Read NICS into a dictionary. - - Returns - ------- - dict - Contains the structure information as well as the nucleus-independent - chemical shift represented on a grid in the unit cell. - """ - result = { - "structure": self._structure().to_dict(), - "nics": self.to_numpy(), - **self._get_method_and_positions(), - } - return result - - def to_database(self) -> dict: - method = "grid" if self._data_is_on_grid else "positions" - return Nics_DB(method=method) - - def to_numpy(self, selection: Optional[str] = None): - selected_data = self._read_selected_data(selection) - return np.squeeze(list(selected_data.values())) - - def to_view( - self, - selection: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - **user_options, - ): - self._raise_error_if_used_in_points_mode() - selection = selection or _DEFAULT_SELECTION - viewer = self._structure().to_view(supercell) - viewer.grid_scalars = [ - self._make_grid_quantity(*item, user_options) - for item in self._read_selected_data(selection).items() - ] - return viewer - - def to_contour( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ): - self._raise_error_if_used_in_points_mode() - selection = selection or _DEFAULT_SELECTION - cut, fraction = slicing.get_cut(a, b, c) - plane = slicing.plane(self._structure().lattice_vectors(), cut, normal) - contour_plots = [ - self._make_contour(*item, plane, fraction, supercell) - for item in self._read_selected_data(selection).items() - ] - return graph.Graph(contour_plots) - - def _structure(self): - return StructureHandler.from_data(self._raw_nics.structure) - - @property - def _data_is_on_grid(self): - return check.is_none(self._raw_nics.positions) - - def _read_selected_data(self, selection: Optional[str]): - if self._data_is_on_grid: - nics_data = np.array(self._raw_nics.nics_grid).T - else: - nics_data = np.array(self._raw_nics.nics_points) - nics_data = nics_data.reshape((len(nics_data), 9)) - if selection is None: - new_shape = (*nics_data.shape[:-1], 3, 3) - return {None: nics_data.reshape(new_shape)} - tree = select.Tree.from_selection(selection) - maps = {nics_data.ndim - 1: self._init_directions_dict()} - selector = index.Selector(maps, nics_data, reduction=_TensorReduction) - return { - selector.label(selection): selector[selection] - for selection in tree.selections() - } - - @staticmethod - def _init_directions_dict(): - return { - "isotropic": [0, 4, 8], - "xx": 0, - "xy": 1, - "xz": 2, - "yx": 3, - "yy": 4, - "yz": 5, - "zx": 6, - "zy": 7, - "zz": 8, - "11": slice(None), - "22": slice(None), - "33": slice(None), - "span": slice(None), - "skew": slice(None), - "anisotropy": slice(None), - "asymmetry": slice(None), - } - - def _get_method_and_positions(self): - if self._data_is_on_grid: - return {"method": "grid"} - else: - return {"method": "positions", "positions": self._raw_nics.positions[:].T} - - def _grid_to_string(self): - grid = self._raw_nics.nics_grid.shape[1:] - return f""" grid: {grid[2]}, {grid[1]}, {grid[0]} - tensor shape: 3x3""" - - def _points_to_string(self): - positions = self._raw_nics.positions[:].T - tensors = self.to_numpy() - return "\n\n".join(self._format_nics(*item) for item in zip(positions, tensors)) - - def _format_nics(self, position, tensor): - position_string = " ".join(f"{x:10.6f}" for x in position) - newline_with_indent = "\n " - tensor = np.round(tensor, 14) - tensor_string = newline_with_indent.join( - " ".join(f"{x:+.6e}" for x in column) for column in tensor - ) - return f"""\ - NICS at {position_string}: | - {tensor_string}""" - - def _make_grid_quantity(self, key, quantity, user_options): - return view.GridQuantity( - quantity=quantity[np.newaxis], - label=f"{key} NICS", - isosurfaces=self._isosurfaces(**user_options), - ) - - def _isosurfaces(self, isolevel=1.0, opacity=0.6): - return [ - view.Isosurface(isolevel, _config.VASP_COLORS["blue"], opacity), - view.Isosurface(-isolevel, _config.VASP_COLORS["red"], opacity), - ] - - def _make_contour(self, key, data, plane, fraction, supercell): - grid_scalar = slicing.grid_scalar(data, plane, fraction) - label = f"{key} NICS contour ({plane.cut})" - contour_plot = graph.Contour(grid_scalar, plane, label, isolevels=True) - if supercell is not None: - contour_plot.supercell = np.ones(2, dtype=np.int_) * supercell - return contour_plot - - def _raise_error_if_used_in_points_mode(self): - if self._data_is_on_grid: - return - raise exception.IncorrectUsage( - "You set LNICSALL = .FALSE. in the INCAR file. This mode is incompatible with the plotting routines." - ) - - -@quantity("nics") -class Nics(view.Mixin): - """This class accesses information on the nucleus-independent chemical shift (NICS). - - Examples - -------- - - First, we create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - See some basic information about NICS by printing the object: - - >>> print(calculation.nics) - nucleus-independent chemical shift:... - - For your own postprocessing, you can read the band data into a Python - dictionary: - - >>> calculation.nics.read() - {'structure': {...}, 'nics': array([[[[[...]]]]]...), 'method': ...} - - You can also obtain the NICS as a numpy array directly: - - >>> calculation.nics.to_numpy() - array([[[[[...]]]]]...) - - You can also visualize a 3d isosurface of the chemical shift: - - >>> calculation.nics.plot() - View(elements=array([[...]]...), lattice_vectors=array([[[...]]]...), positions=array([[[...]]]...), grid_scalars=[GridQuantity(quantity=array([[[[...]]]]...), label='isotropic NICS', isosurfaces=[Isosurface(...)])], ...) - - Alternatively, you can visualize a contour plot of the chemical shift in a plane: - - >>> calculation.nics.to_contour(c=0) - Graph(series=[Contour(data=array([[...]]...), ..., cut='c', ...)], ...) - - Please check the documentation of each of these methods for more details on how to - use them and which options they provide. - """ - - def __init__(self, source, quantity_name="nics"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_nics): - return cls(source=DataSource(raw_nics)) - - def _handler_factory(self, raw): - return NicsHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - NicsHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Read NICS into a dictionary. - - Returns - ------- - dict - Contains the structure information as well as the nucleus-independent - chemical shift represented on a grid in the unit cell. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - NicsHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def to_numpy(self, selection: Optional[str] = None): - """Convert NICS to a numpy array. - - The resulting shape will be the NICS grid data with respect to the selection. - - Parameters - ---------- - selection : str or None - The tensor element(s) to extract. - Can be None (in which case the whole tensor is returned), isotropic, or - one of "xx", "xy", ... - - Returns - ------- - np.ndarray - All components of NICS. - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - NicsHandler.to_numpy, - ) - - def to_view( - self, - selection: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - **user_options, - ): - """Plot the selected chemical shift as a 3d isosurface within the structure. - - Parameters - ---------- - selection : str or None - Axis along which to plot. - Can be one of "xx", "xy", ... - Can also be "isotropic" to plot the trace. - If selection is None, it defaults to "isotropic". - supercell : int or np.ndarray - If present the data is replicated the specified number of times along each - direction. - user_options - Further arguments with keyword that get directly passed on to the - visualizer. Most importantly, you can set isolevel to adjust the - value at which the isosurface is drawn. - - Returns - ------- - View - Visualize an isosurface of the selected chemical shift within the 3d - structure. - - Examples - -------- - >>> from py4vasp import calculation - - Plot the isotropic chemical shift as a 3d isosurface. - - >>> calculation.nics.plot() - - Plot the chemical shift with "xx" selection as a 3d isosurface. - - >>> calculation.nics.plot(selection="xx") - - Plot the isotropic chemical shift with specified isolevel as a 3d isosurface. - - >>> calculation.nics.plot(isolevel=0.6) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - NicsHandler.to_view, - supercell=supercell, - **user_options, - ) - - @documentation.format(plane=slicing.PLANE, parameters=slicing.PARAMETERS) - def to_contour( - self, - selection: Optional[str] = None, - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ): - """Generate a contour plot of chemical shift. - - {plane} - - Parameters - ---------- - {parameters} - selection : str or None - Axis along which to plot. - Can be one of "xx", "xy", ... - Can also be "isotropic" to plot the trace. - If selection is None, it defaults to "isotropic". - supercell : int or np.ndarray - If present the data is replicated the specified number of times along each - direction. - - Returns - ------- - graph - A chemical shift plot in the plane spanned by the 2 remaining lattice - vectors. - - Examples - -------- - >>> from py4vasp import calculation - - Cut a plane through the isotropic chemical shift at the origin of the third - lattice vector. - - >>> calculation.nics.to_contour(c=0) - - Replicate a plane in the middle of the second lattice vector 2 times in each - direction. - - >>> calculation.nics.to_contour(b=0.5, supercell=2) - - Take a slice of the chemical shift with "xy" selection along the first lattice - vector and rotate it such that the plane normal aligns with the x axis. - - >>> calculation.nics.to_contour(a=0.3, selection=0.3, normal="x") - - Cut a plane through the isotropic chemical shift at the origin of the third - lattice vector, then show isosurface level values along contour lines. - - >>> plot = calculation.nics.to_contour(c=0, selection=0.3, normal="x") - >>> plot.series[0].show_contour_values = True - >>> plot.show() - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - NicsHandler.to_contour, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - NicsHandler.from_data, - NicsHandler.to_database, - ) - - -class _TensorReduction(index.Reduction): - def __init__(self, keys): - keys_using_average = "isotropic xx xy xz yx yy yz zx zy zz" - self._use_average = keys[-1] in keys_using_average - self._selection = keys[-1] - - def __call__(self, array, axis): - if self._use_average: - return np.average(array, axis=axis) - else: - return self._reduce(array, axis) - - def _reduce(self, array, axis): - array = array.reshape((*array.shape[:-1], 3, 3)) - symmetric_array = 0.5 * (array + np.moveaxis(array, -2, -1)) - eigenvalues = np.linalg.eigvalsh(array) - if self._selection == "11": - return eigenvalues[..., 2] - if self._selection == "22": - return eigenvalues[..., 1] - if self._selection == "33": - return eigenvalues[..., 0] - if self._selection == "span": - return eigenvalues[..., 2] - eigenvalues[..., 0] - if self._selection == "skew": - span = eigenvalues[..., 2] - eigenvalues[..., 0] - return (3 * eigenvalues[..., 1] - np.sum(eigenvalues, axis=-1)) / span - if self._selection in ("anisotropy", "asymmetry"): - return self._haeberlen_mehring(eigenvalues)[self._selection] - message = f"The reduction for selection '{self._selection}' is not implemented." - raise exception.NotImplemented(message) - - def _haeberlen_mehring(self, eigenvalues): - delta_iso = np.average(eigenvalues, axis=-1) - mask = delta_iso < eigenvalues[..., 1] - delta_xx = np.where(mask, eigenvalues[..., 2], eigenvalues[..., 0]) - delta_zz = np.where(mask, eigenvalues[..., 0], eigenvalues[..., 2]) - anisotropy = delta_zz - delta_iso - asymmetry = (eigenvalues[..., 1] - delta_xx) / anisotropy - return {"anisotropy": anisotropy, "asymmetry": asymmetry} +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy +from typing import Optional, Union + +import numpy as np + +from py4vasp import _config, exception +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Nics_DB +from py4vasp._third_party import graph, view +from py4vasp._util import check, documentation, import_, index, select, slicing + +pretty = import_.optional("IPython.lib.pretty") + +_DEFAULT_SELECTION: str = "isotropic" + + +class NicsHandler: + """Handler for NICS data — performs all data access and transformation.""" + + def __init__(self, raw_nics: raw.Nics): + self._raw_nics = raw_nics + + @classmethod + def from_data(cls, raw_nics: raw.Nics) -> "NicsHandler": + return cls(raw_nics) + + def __str__(self) -> str: + raw_stoichiometry = self._raw_nics.structure.stoichiometry + stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry) + if self._data_is_on_grid: + data_string = self._grid_to_string() + else: + data_string = self._points_to_string() + return f"""\ +nucleus-independent chemical shift: + structure: {pretty.pretty(stoichiometry)} +{data_string}""" + + def to_dict(self) -> dict: + """Read NICS into a dictionary. + + Returns + ------- + dict + Contains the structure information as well as the nucleus-independent + chemical shift represented on a grid in the unit cell. + """ + result = { + "structure": self._structure().to_dict(), + "nics": self.to_numpy(), + **self._get_method_and_positions(), + } + return result + + def to_database(self) -> dict: + method = "grid" if self._data_is_on_grid else "positions" + return Nics_DB(method=method) + + def to_numpy(self, selection: Optional[str] = None): + selected_data = self._read_selected_data(selection) + return np.squeeze(list(selected_data.values())) + + def to_view( + self, + selection: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + **user_options, + ): + self._raise_error_if_used_in_points_mode() + selection = selection or _DEFAULT_SELECTION + viewer = self._structure().to_view(supercell) + viewer.grid_scalars = [ + self._make_grid_quantity(*item, user_options) + for item in self._read_selected_data(selection).items() + ] + return viewer + + def to_contour( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ): + self._raise_error_if_used_in_points_mode() + selection = selection or _DEFAULT_SELECTION + cut, fraction = slicing.get_cut(a, b, c) + plane = slicing.plane(self._structure().lattice_vectors(), cut, normal) + contour_plots = [ + self._make_contour(*item, plane, fraction, supercell) + for item in self._read_selected_data(selection).items() + ] + return graph.Graph(contour_plots) + + def _structure(self): + return StructureHandler.from_data(self._raw_nics.structure) + + @property + def _data_is_on_grid(self): + return check.is_none(self._raw_nics.positions) + + def _read_selected_data(self, selection: Optional[str]): + if self._data_is_on_grid: + nics_data = np.array(self._raw_nics.nics_grid).T + else: + nics_data = np.array(self._raw_nics.nics_points) + nics_data = nics_data.reshape((len(nics_data), 9)) + if selection is None: + new_shape = (*nics_data.shape[:-1], 3, 3) + return {None: nics_data.reshape(new_shape)} + tree = select.Tree.from_selection(selection) + maps = {nics_data.ndim - 1: self._init_directions_dict()} + selector = index.Selector(maps, nics_data, reduction=_TensorReduction) + return { + selector.label(selection): selector[selection] + for selection in tree.selections() + } + + @staticmethod + def _init_directions_dict(): + return { + "isotropic": [0, 4, 8], + "xx": 0, + "xy": 1, + "xz": 2, + "yx": 3, + "yy": 4, + "yz": 5, + "zx": 6, + "zy": 7, + "zz": 8, + "11": slice(None), + "22": slice(None), + "33": slice(None), + "span": slice(None), + "skew": slice(None), + "anisotropy": slice(None), + "asymmetry": slice(None), + } + + def _get_method_and_positions(self): + if self._data_is_on_grid: + return {"method": "grid"} + else: + return {"method": "positions", "positions": self._raw_nics.positions[:].T} + + def _grid_to_string(self): + grid = self._raw_nics.nics_grid.shape[1:] + return f""" grid: {grid[2]}, {grid[1]}, {grid[0]} + tensor shape: 3x3""" + + def _points_to_string(self): + positions = self._raw_nics.positions[:].T + tensors = self.to_numpy() + return "\n\n".join(self._format_nics(*item) for item in zip(positions, tensors)) + + def _format_nics(self, position, tensor): + position_string = " ".join(f"{x:10.6f}" for x in position) + newline_with_indent = "\n " + tensor = np.round(tensor, 14) + tensor_string = newline_with_indent.join( + " ".join(f"{x:+.6e}" for x in column) for column in tensor + ) + return f"""\ + NICS at {position_string}: | + {tensor_string}""" + + def _make_grid_quantity(self, key, quantity, user_options): + return view.GridQuantity( + quantity=quantity[np.newaxis], + label=f"{key} NICS", + isosurfaces=self._isosurfaces(**user_options), + ) + + def _isosurfaces(self, isolevel=1.0, opacity=0.6): + return [ + view.Isosurface(isolevel, _config.VASP_COLORS["blue"], opacity), + view.Isosurface(-isolevel, _config.VASP_COLORS["red"], opacity), + ] + + def _make_contour(self, key, data, plane, fraction, supercell): + grid_scalar = slicing.grid_scalar(data, plane, fraction) + label = f"{key} NICS contour ({plane.cut})" + contour_plot = graph.Contour(grid_scalar, plane, label, isolevels=True) + if supercell is not None: + contour_plot.supercell = np.ones(2, dtype=np.int_) * supercell + return contour_plot + + def _raise_error_if_used_in_points_mode(self): + if self._data_is_on_grid: + return + raise exception.IncorrectUsage( + "You set LNICSALL = .FALSE. in the INCAR file. This mode is incompatible with the plotting routines." + ) + + +@quantity("nics") +class Nics(view.Mixin): + """This class accesses information on the nucleus-independent chemical shift (NICS). + + Examples + -------- + + First, we create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + See some basic information about NICS by printing the object: + + >>> print(calculation.nics) + nucleus-independent chemical shift:... + + For your own postprocessing, you can read the band data into a Python + dictionary: + + >>> calculation.nics.read() + {'structure': {...}, 'nics': array([[[[[...]]]]]...), 'method': ...} + + You can also obtain the NICS as a numpy array directly: + + >>> calculation.nics.to_numpy() + array([[[[[...]]]]]...) + + You can also visualize a 3d isosurface of the chemical shift: + + >>> calculation.nics.plot() + View(elements=array([[...]]...), lattice_vectors=array([[[...]]]...), positions=array([[[...]]]...), grid_scalars=[GridQuantity(quantity=array([[[[...]]]]...), label='isotropic NICS', isosurfaces=[Isosurface(...)])], ...) + + Alternatively, you can visualize a contour plot of the chemical shift in a plane: + + >>> calculation.nics.to_contour(c=0) + Graph(series=[Contour(data=array([[...]]...), ..., cut='c', ...)], ...) + + Please check the documentation of each of these methods for more details on how to + use them and which options they provide. + """ + + def __init__(self, source, quantity_name="nics"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_nics): + return cls(source=DataSource(raw_nics)) + + def _handler_factory(self, raw): + return NicsHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + NicsHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Read NICS into a dictionary. + + Returns + ------- + dict + Contains the structure information as well as the nucleus-independent + chemical shift represented on a grid in the unit cell. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + NicsHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def to_numpy(self, selection: Optional[str] = None): + """Convert NICS to a numpy array. + + The resulting shape will be the NICS grid data with respect to the selection. + + Parameters + ---------- + selection : str or None + The tensor element(s) to extract. + Can be None (in which case the whole tensor is returned), isotropic, or + one of "xx", "xy", ... + + Returns + ------- + np.ndarray + All components of NICS. + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + NicsHandler.to_numpy, + ) + + def to_view( + self, + selection: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + **user_options, + ): + """Plot the selected chemical shift as a 3d isosurface within the structure. + + Parameters + ---------- + selection : str or None + Axis along which to plot. + Can be one of "xx", "xy", ... + Can also be "isotropic" to plot the trace. + If selection is None, it defaults to "isotropic". + supercell : int or np.ndarray + If present the data is replicated the specified number of times along each + direction. + user_options + Further arguments with keyword that get directly passed on to the + visualizer. Most importantly, you can set isolevel to adjust the + value at which the isosurface is drawn. + + Returns + ------- + View + Visualize an isosurface of the selected chemical shift within the 3d + structure. + + Examples + -------- + >>> from py4vasp import calculation + + Plot the isotropic chemical shift as a 3d isosurface. + + >>> calculation.nics.plot() + + Plot the chemical shift with "xx" selection as a 3d isosurface. + + >>> calculation.nics.plot(selection="xx") + + Plot the isotropic chemical shift with specified isolevel as a 3d isosurface. + + >>> calculation.nics.plot(isolevel=0.6) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + NicsHandler.to_view, + supercell=supercell, + **user_options, + ) + + @documentation.format(plane=slicing.PLANE, parameters=slicing.PARAMETERS) + def to_contour( + self, + selection: Optional[str] = None, + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ): + """Generate a contour plot of chemical shift. + + {plane} + + Parameters + ---------- + {parameters} + selection : str or None + Axis along which to plot. + Can be one of "xx", "xy", ... + Can also be "isotropic" to plot the trace. + If selection is None, it defaults to "isotropic". + supercell : int or np.ndarray + If present the data is replicated the specified number of times along each + direction. + + Returns + ------- + graph + A chemical shift plot in the plane spanned by the 2 remaining lattice + vectors. + + Examples + -------- + >>> from py4vasp import calculation + + Cut a plane through the isotropic chemical shift at the origin of the third + lattice vector. + + >>> calculation.nics.to_contour(c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calculation.nics.to_contour(b=0.5, supercell=2) + + Take a slice of the chemical shift with "xy" selection along the first lattice + vector and rotate it such that the plane normal aligns with the x axis. + + >>> calculation.nics.to_contour(a=0.3, selection=0.3, normal="x") + + Cut a plane through the isotropic chemical shift at the origin of the third + lattice vector, then show isosurface level values along contour lines. + + >>> plot = calculation.nics.to_contour(c=0, selection=0.3, normal="x") + >>> plot.series[0].show_contour_values = True + >>> plot.show() + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + NicsHandler.to_contour, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + NicsHandler.from_data, + NicsHandler.to_database, + ) + + +class _TensorReduction(index.Reduction): + def __init__(self, keys): + keys_using_average = "isotropic xx xy xz yx yy yz zx zy zz" + self._use_average = keys[-1] in keys_using_average + self._selection = keys[-1] + + def __call__(self, array, axis): + if self._use_average: + return np.average(array, axis=axis) + else: + return self._reduce(array, axis) + + def _reduce(self, array, axis): + array = array.reshape((*array.shape[:-1], 3, 3)) + symmetric_array = 0.5 * (array + np.moveaxis(array, -2, -1)) + eigenvalues = np.linalg.eigvalsh(array) + if self._selection == "11": + return eigenvalues[..., 2] + if self._selection == "22": + return eigenvalues[..., 1] + if self._selection == "33": + return eigenvalues[..., 0] + if self._selection == "span": + return eigenvalues[..., 2] - eigenvalues[..., 0] + if self._selection == "skew": + span = eigenvalues[..., 2] - eigenvalues[..., 0] + return (3 * eigenvalues[..., 1] - np.sum(eigenvalues, axis=-1)) / span + if self._selection in ("anisotropy", "asymmetry"): + return self._haeberlen_mehring(eigenvalues)[self._selection] + message = f"The reduction for selection '{self._selection}' is not implemented." + raise exception.NotImplemented(message) + + def _haeberlen_mehring(self, eigenvalues): + delta_iso = np.average(eigenvalues, axis=-1) + mask = delta_iso < eigenvalues[..., 1] + delta_xx = np.where(mask, eigenvalues[..., 2], eigenvalues[..., 0]) + delta_zz = np.where(mask, eigenvalues[..., 0], eigenvalues[..., 2]) + anisotropy = delta_zz - delta_iso + asymmetry = (eigenvalues[..., 1] - delta_xx) / anisotropy + return {"anisotropy": anisotropy, "asymmetry": asymmetry} diff --git a/src/py4vasp/_calculation/pair_correlation.py b/src/py4vasp/_calculation/pair_correlation.py index fbcee11c6..ffdf2911e 100644 --- a/src/py4vasp/_calculation/pair_correlation.py +++ b/src/py4vasp/_calculation/pair_correlation.py @@ -1,253 +1,253 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy - -from py4vasp import raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import PairCorrelation_DB -from py4vasp._third_party import graph -from py4vasp._util import check, convert, documentation, index, select - - -def _selection_string(default): - return f"""\ -selection : str - String specifying which pair-correlation functions are used. Select - 'total' for the total pair-correlation function or the name of any - two ion types (e.g. 'Sr~Ti') for a specific pair-correlation function. - When no selection is given, {default}. Separate - distinct labels by commas or whitespace. For a complete list of all - possible selections, please use - - >>> calculation.pair_correlation.labels() -""" - - -class PairCorrelationHandler: - """Handler for pair-correlation data — all data access and transformation.""" - - def __init__(self, raw_pair_correlation: raw.PairCorrelation, steps=None): - self._raw_data = raw_pair_correlation - self._steps = steps - - @classmethod - def from_data( - cls, raw_pair_correlation: raw.PairCorrelation, steps=None - ) -> "PairCorrelationHandler": - return cls(raw_pair_correlation, steps) - - def to_dict(self, selection=None) -> dict: - """Read the pair-correlation function and store it in a dictionary.""" - selection = self._default_selection_if_none(selection) - return { - "distances": self._raw_data.distances[:], - **self._read_data(selection), - } - - def to_graph(self, selection="total") -> graph.Graph: - """Plot selected pair-correlation functions.""" - series = self._make_series(self.to_dict(selection)) - return graph.Graph(series, xlabel="Distance (Å)", ylabel="Pair correlation") - - def labels(self) -> tuple: - """Return all possible labels for the selection string.""" - return tuple(convert.text_to_string(label) for label in self._raw_data.labels) - - def to_database(self) -> dict: - """Serialize pair-correlation data for database storage.""" - distance_min, distance_max = None, None - if not check.is_none(self._raw_data.distances): - distance_min = float(self._raw_data.distances[0]) - distance_max = float(self._raw_data.distances[-1]) - return PairCorrelation_DB(distance_min=distance_min, distance_max=distance_max) - - @property - def _steps_or_last(self): - return -1 if self._steps is None else self._steps - - def _default_selection_if_none(self, selection): - return selection or ",".join(self.labels()) - - def _read_data(self, selection): - map_ = {1: self._init_pair_correlation_dict()} - selector = index.Selector(map_, self._raw_data.function) - tree = select.Tree.from_selection(selection) - return { - selector.label(selection): selector[selection][self._steps_or_last] - for selection in tree.selections() - } - - def _init_pair_correlation_dict(self): - return {label: i for i, label in enumerate(self.labels())} - - def _make_series(self, selected_data): - distances = selected_data["distances"] - return [ - graph.Series(x=distances, y=data, label=label) - for label, data in selected_data.items() - if label != "distances" - ] - - -@quantity("pair_correlation") -@documentation.format(examples=slice_.examples("pair_correlation", step="block")) -class PairCorrelation(graph.Mixin): - """The pair-correlation function measures the distribution of atoms. - - A pair-correlation function is a statistical measure to describe the spatial - distribution of atoms within a system. Specifically, the pair correlation - function quantifies the probability density of finding two particles at specific - separation distances. This function is helpful in the study of liquids and solids - because it acts as a fingerprint of the system that can be compared to - X-ray or neutron scattering experiments. Another use case is the detection - of specific phases. - - Use this class to inspect the pair-correlation function computed by VASP for - all pairs of ionic types. You can control how often VASP samples the pair - correlation function with the :tag:`NBLOCK` tag. If you want to split your - trajectory into multiple subsets include the tag :tag:`KBLOCK` in your INCAR - file. - - {examples} - """ - - def __init__(self, source, quantity_name: str = "pair_correlation", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_pair_correlation: raw.PairCorrelation): - """Create a PairCorrelation dispatcher from raw data.""" - return cls(source=DataSource(raw_pair_correlation)) - - @property - def path(self): - """Path used for file-export methods.""" - return self._path - - def __getitem__(self, steps) -> "PairCorrelation": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw_data): - return PairCorrelationHandler.from_data(raw_data, steps=self._steps) - - @documentation.format( - selection=_selection_string("all possibilities are read"), - examples=slice_.examples("pair_correlation", "to_dict", "block"), - ) - def read(self, selection=None) -> dict: - """Read the pair-correlation function and store it in a dictionary. - - Parameters - ---------- - {selection} - - Returns - ------- - dict - Contains the labels corresponding to the selection and the associated - pair-correlation function for every selected block. Furthermore, the - dictionary contains the distances at which the pair-correlation functions - are evaluated. - - {examples} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PairCorrelationHandler.to_dict, - ) - - @documentation.format( - selection=_selection_string("all possibilities are read"), - examples=slice_.examples("pair_correlation", "to_dict", "block"), - ) - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read(selection=selection) - - @documentation.format( - selection=_selection_string("the total pair correlation is used"), - examples=slice_.examples("pair_correlation", "to_graph", "block"), - ) - def to_graph(self, selection="total") -> graph.Graph: - """Plot selected pair-correlation functions. - - Parameters - ---------- - {selection} - - Returns - ------- - Graph - The graph plots the pair-correlation function for all selected blocks - and ion pairs. - - {examples} - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PairCorrelationHandler.to_graph, - ) - - def labels(self) -> tuple: - """Return all possible labels for the selection string.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - PairCorrelationHandler.labels, - ) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PairCorrelationHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PairCorrelationHandler.from_data, - PairCorrelationHandler.to_database, - ) - - -def _selection_string(default): - return f"""\ -selection : str - String specifying which pair-correlation functions are used. Select - 'total' for the total pair-correlation function or the name of any - two ion types (e.g. 'Sr~Ti') for a specific pair-correlation function. - When no selection is given, {default}. Separate - distinct labels by commas or whitespace. For a complete list of all - possible selections, please use - - >>> calculation.pair_correlation.labels() -""" +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy + +from py4vasp import raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import PairCorrelation_DB +from py4vasp._third_party import graph +from py4vasp._util import check, convert, documentation, index, select + + +def _selection_string(default): + return f"""\ +selection : str + String specifying which pair-correlation functions are used. Select + 'total' for the total pair-correlation function or the name of any + two ion types (e.g. 'Sr~Ti') for a specific pair-correlation function. + When no selection is given, {default}. Separate + distinct labels by commas or whitespace. For a complete list of all + possible selections, please use + + >>> calculation.pair_correlation.labels() +""" + + +class PairCorrelationHandler: + """Handler for pair-correlation data — all data access and transformation.""" + + def __init__(self, raw_pair_correlation: raw.PairCorrelation, steps=None): + self._raw_data = raw_pair_correlation + self._steps = steps + + @classmethod + def from_data( + cls, raw_pair_correlation: raw.PairCorrelation, steps=None + ) -> "PairCorrelationHandler": + return cls(raw_pair_correlation, steps) + + def to_dict(self, selection=None) -> dict: + """Read the pair-correlation function and store it in a dictionary.""" + selection = self._default_selection_if_none(selection) + return { + "distances": self._raw_data.distances[:], + **self._read_data(selection), + } + + def to_graph(self, selection="total") -> graph.Graph: + """Plot selected pair-correlation functions.""" + series = self._make_series(self.to_dict(selection)) + return graph.Graph(series, xlabel="Distance (Å)", ylabel="Pair correlation") + + def labels(self) -> tuple: + """Return all possible labels for the selection string.""" + return tuple(convert.text_to_string(label) for label in self._raw_data.labels) + + def to_database(self) -> dict: + """Serialize pair-correlation data for database storage.""" + distance_min, distance_max = None, None + if not check.is_none(self._raw_data.distances): + distance_min = float(self._raw_data.distances[0]) + distance_max = float(self._raw_data.distances[-1]) + return PairCorrelation_DB(distance_min=distance_min, distance_max=distance_max) + + @property + def _steps_or_last(self): + return -1 if self._steps is None else self._steps + + def _default_selection_if_none(self, selection): + return selection or ",".join(self.labels()) + + def _read_data(self, selection): + map_ = {1: self._init_pair_correlation_dict()} + selector = index.Selector(map_, self._raw_data.function) + tree = select.Tree.from_selection(selection) + return { + selector.label(selection): selector[selection][self._steps_or_last] + for selection in tree.selections() + } + + def _init_pair_correlation_dict(self): + return {label: i for i, label in enumerate(self.labels())} + + def _make_series(self, selected_data): + distances = selected_data["distances"] + return [ + graph.Series(x=distances, y=data, label=label) + for label, data in selected_data.items() + if label != "distances" + ] + + +@quantity("pair_correlation") +@documentation.format(examples=slice_.examples("pair_correlation", step="block")) +class PairCorrelation(graph.Mixin): + """The pair-correlation function measures the distribution of atoms. + + A pair-correlation function is a statistical measure to describe the spatial + distribution of atoms within a system. Specifically, the pair correlation + function quantifies the probability density of finding two particles at specific + separation distances. This function is helpful in the study of liquids and solids + because it acts as a fingerprint of the system that can be compared to + X-ray or neutron scattering experiments. Another use case is the detection + of specific phases. + + Use this class to inspect the pair-correlation function computed by VASP for + all pairs of ionic types. You can control how often VASP samples the pair + correlation function with the :tag:`NBLOCK` tag. If you want to split your + trajectory into multiple subsets include the tag :tag:`KBLOCK` in your INCAR + file. + + {examples} + """ + + def __init__(self, source, quantity_name: str = "pair_correlation", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_pair_correlation: raw.PairCorrelation): + """Create a PairCorrelation dispatcher from raw data.""" + return cls(source=DataSource(raw_pair_correlation)) + + @property + def path(self): + """Path used for file-export methods.""" + return self._path + + def __getitem__(self, steps) -> "PairCorrelation": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw_data): + return PairCorrelationHandler.from_data(raw_data, steps=self._steps) + + @documentation.format( + selection=_selection_string("all possibilities are read"), + examples=slice_.examples("pair_correlation", "to_dict", "block"), + ) + def read(self, selection=None) -> dict: + """Read the pair-correlation function and store it in a dictionary. + + Parameters + ---------- + {selection} + + Returns + ------- + dict + Contains the labels corresponding to the selection and the associated + pair-correlation function for every selected block. Furthermore, the + dictionary contains the distances at which the pair-correlation functions + are evaluated. + + {examples} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PairCorrelationHandler.to_dict, + ) + + @documentation.format( + selection=_selection_string("all possibilities are read"), + examples=slice_.examples("pair_correlation", "to_dict", "block"), + ) + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read(selection=selection) + + @documentation.format( + selection=_selection_string("the total pair correlation is used"), + examples=slice_.examples("pair_correlation", "to_graph", "block"), + ) + def to_graph(self, selection="total") -> graph.Graph: + """Plot selected pair-correlation functions. + + Parameters + ---------- + {selection} + + Returns + ------- + Graph + The graph plots the pair-correlation function for all selected blocks + and ion pairs. + + {examples} + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PairCorrelationHandler.to_graph, + ) + + def labels(self) -> tuple: + """Return all possible labels for the selection string.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + PairCorrelationHandler.labels, + ) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PairCorrelationHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PairCorrelationHandler.from_data, + PairCorrelationHandler.to_database, + ) + + +def _selection_string(default): + return f"""\ +selection : str + String specifying which pair-correlation functions are used. Select + 'total' for the total pair-correlation function or the name of any + two ion types (e.g. 'Sr~Ti') for a specific pair-correlation function. + When no selection is given, {default}. Separate + distinct labels by commas or whitespace. For a complete list of all + possible selections, please use + + >>> calculation.pair_correlation.labels() +""" diff --git a/src/py4vasp/_calculation/phonon_band.py b/src/py4vasp/_calculation/phonon_band.py index 99f2f1308..c7888911a 100644 --- a/src/py4vasp/_calculation/phonon_band.py +++ b/src/py4vasp/_calculation/phonon_band.py @@ -1,206 +1,206 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import pathlib - -import numpy as np - -from py4vasp import raw -from py4vasp._calculation import phonon -from py4vasp._calculation._dispersion import DispersionHandler -from py4vasp._calculation._stoichiometry import StoichiometryHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._third_party import graph -from py4vasp._util import convert, documentation, index, select - - -class PhononBandHandler: - """Handler for phonon band structure data.""" - - def __init__(self, raw_phonon_band: raw.PhononBand): - self._raw_phonon_band = raw_phonon_band - - @classmethod - def from_data(cls, raw_phonon_band: raw.PhononBand) -> "PhononBandHandler": - return cls(raw_phonon_band) - - def __str__(self) -> str: - return f"""phonon band data: - {self._raw_phonon_band.dispersion.eigenvalues.shape[0]} q-points - {self._raw_phonon_band.dispersion.eigenvalues.shape[1]} modes - {self._stoichiometry()}""" - - def to_dict(self) -> dict: - dispersion = self._dispersion().to_dict() - return { - "qpoint_distances": dispersion["kpoint_distances"], - "qpoint_labels": dispersion.get("kpoint_labels"), - "bands": dispersion["eigenvalues"], - "modes": self._modes(), - } - - def to_database(self) -> dict: - return {} - - def to_graph(self, selection=None, width=1.0) -> graph.Graph: - projections = self._projections(selection, width) - g = self._dispersion().plot(projections) - g.ylabel = "ω (THz)" - return g - - def selections(self) -> dict: - atoms = self._init_atom_dict().keys() - return { - "atom": sorted(atoms, key=self._sort_key), - "direction": ["x", "y", "z"], - } - - def _dispersion(self) -> DispersionHandler: - return DispersionHandler.from_data(self._raw_phonon_band.dispersion) - - def _stoichiometry(self) -> StoichiometryHandler: - return StoichiometryHandler.from_data(self._raw_phonon_band.stoichiometry) - - def _modes(self) -> np.ndarray: - return convert.to_complex(self._raw_phonon_band.eigenvectors[:]) - - def _projections(self, selection, width): - if not selection: - return None - maps = {2: self._init_atom_dict(), 3: self._init_direction_dict()} - selector = index.Selector(maps, np.abs(self._modes()), use_number_labels=True) - tree = select.Tree.from_selection(selection) - return {selector.label(sel): width * selector[sel] for sel in tree.selections()} - - def _init_atom_dict(self) -> dict: - return { - key: value.indices - for key, value in self._stoichiometry().read().items() - if key != select.all - } - - def _init_direction_dict(self) -> dict: - return { - "x": slice(0, 1), - "y": slice(1, 2), - "z": slice(2, 3), - } - - def _sort_key(self, key) -> bool: - return key.isdecimal() - - -@quantity("band", group="phonon") -class PhononBand(graph.Mixin): - """The phonon band structure contains the **q**-resolved phonon eigenvalues. - - The phonon band structure is a graphical representation of the phonons. It - illustrates the relationship between the frequency of modes and their corresponding - wave vectors in the Brillouin zone. Each line or branch in the band structure - represents a specific phonon, and the slope of these branches provides information - about their velocity. - - The phonon band structure includes the dispersion relations of phonons, which reveal - how vibrational frequencies vary with direction in the crystal lattice. The presence - of band gaps or band crossings indicates the material's ability to conduct or - insulate heat. Additionally, the branches near the high-symmetry points in the - Brillouin zone offer insights into the material's anharmonicity and thermal - conductivity. Furthermore, phonons with imaginary frequencies indicate the presence - of a structural instability. - """ - - def __init__(self, source, quantity_name: str = "phonon_band"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_phonon_band: raw.PhononBand) -> "PhononBand": - return cls(source=DataSource(raw_phonon_band)) - - def _handler_factory(self, raw): - return PhononBandHandler.from_data(raw) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononBandHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self, selection=None) -> dict: - """Read the phonon band structure into a dictionary. - - Returns - ------- - dict - Contains the **q**-point path for plotting phonon band structures and - the phonon bands. In addition the phonon modes are returned. - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononBandHandler.to_dict, - ) - - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read(selection=selection) - - @documentation.format(selection=phonon.selection_doc) - def to_graph(self, selection: str | None = None, width: float = 1.0) -> graph.Graph: - """Generate a graph of the phonon bands. - - Parameters - ---------- - {selection} - width : float - Specifies the width illustrating the projections. - - Returns - ------- - Graph - Contains the phonon band structure for all the **q** points. If a - selection is provided, the width of the bands is adjusted according to - the projection. - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononBandHandler.to_graph, - width=width, - ) - - def selections(self, selection=None) -> dict: - """Return atom and direction selections available for projection.""" - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononBandHandler.selections, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PhononBandHandler.from_data, - PhononBandHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import pathlib + +import numpy as np + +from py4vasp import raw +from py4vasp._calculation import phonon +from py4vasp._calculation._dispersion import DispersionHandler +from py4vasp._calculation._stoichiometry import StoichiometryHandler +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._third_party import graph +from py4vasp._util import convert, documentation, index, select + + +class PhononBandHandler: + """Handler for phonon band structure data.""" + + def __init__(self, raw_phonon_band: raw.PhononBand): + self._raw_phonon_band = raw_phonon_band + + @classmethod + def from_data(cls, raw_phonon_band: raw.PhononBand) -> "PhononBandHandler": + return cls(raw_phonon_band) + + def __str__(self) -> str: + return f"""phonon band data: + {self._raw_phonon_band.dispersion.eigenvalues.shape[0]} q-points + {self._raw_phonon_band.dispersion.eigenvalues.shape[1]} modes + {self._stoichiometry()}""" + + def to_dict(self) -> dict: + dispersion = self._dispersion().to_dict() + return { + "qpoint_distances": dispersion["kpoint_distances"], + "qpoint_labels": dispersion.get("kpoint_labels"), + "bands": dispersion["eigenvalues"], + "modes": self._modes(), + } + + def to_database(self) -> dict: + return {} + + def to_graph(self, selection=None, width=1.0) -> graph.Graph: + projections = self._projections(selection, width) + g = self._dispersion().plot(projections) + g.ylabel = "ω (THz)" + return g + + def selections(self) -> dict: + atoms = self._init_atom_dict().keys() + return { + "atom": sorted(atoms, key=self._sort_key), + "direction": ["x", "y", "z"], + } + + def _dispersion(self) -> DispersionHandler: + return DispersionHandler.from_data(self._raw_phonon_band.dispersion) + + def _stoichiometry(self) -> StoichiometryHandler: + return StoichiometryHandler.from_data(self._raw_phonon_band.stoichiometry) + + def _modes(self) -> np.ndarray: + return convert.to_complex(self._raw_phonon_band.eigenvectors[:]) + + def _projections(self, selection, width): + if not selection: + return None + maps = {2: self._init_atom_dict(), 3: self._init_direction_dict()} + selector = index.Selector(maps, np.abs(self._modes()), use_number_labels=True) + tree = select.Tree.from_selection(selection) + return {selector.label(sel): width * selector[sel] for sel in tree.selections()} + + def _init_atom_dict(self) -> dict: + return { + key: value.indices + for key, value in self._stoichiometry().read().items() + if key != select.all + } + + def _init_direction_dict(self) -> dict: + return { + "x": slice(0, 1), + "y": slice(1, 2), + "z": slice(2, 3), + } + + def _sort_key(self, key) -> bool: + return key.isdecimal() + + +@quantity("band", group="phonon") +class PhononBand(graph.Mixin): + """The phonon band structure contains the **q**-resolved phonon eigenvalues. + + The phonon band structure is a graphical representation of the phonons. It + illustrates the relationship between the frequency of modes and their corresponding + wave vectors in the Brillouin zone. Each line or branch in the band structure + represents a specific phonon, and the slope of these branches provides information + about their velocity. + + The phonon band structure includes the dispersion relations of phonons, which reveal + how vibrational frequencies vary with direction in the crystal lattice. The presence + of band gaps or band crossings indicates the material's ability to conduct or + insulate heat. Additionally, the branches near the high-symmetry points in the + Brillouin zone offer insights into the material's anharmonicity and thermal + conductivity. Furthermore, phonons with imaginary frequencies indicate the presence + of a structural instability. + """ + + def __init__(self, source, quantity_name: str = "phonon_band"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_phonon_band: raw.PhononBand) -> "PhononBand": + return cls(source=DataSource(raw_phonon_band)) + + def _handler_factory(self, raw): + return PhononBandHandler.from_data(raw) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononBandHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self, selection=None) -> dict: + """Read the phonon band structure into a dictionary. + + Returns + ------- + dict + Contains the **q**-point path for plotting phonon band structures and + the phonon bands. In addition the phonon modes are returned. + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononBandHandler.to_dict, + ) + + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read(selection=selection) + + @documentation.format(selection=phonon.selection_doc) + def to_graph(self, selection: str | None = None, width: float = 1.0) -> graph.Graph: + """Generate a graph of the phonon bands. + + Parameters + ---------- + {selection} + width : float + Specifies the width illustrating the projections. + + Returns + ------- + Graph + Contains the phonon band structure for all the **q** points. If a + selection is provided, the width of the bands is adjusted according to + the projection. + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononBandHandler.to_graph, + width=width, + ) + + def selections(self, selection=None) -> dict: + """Return atom and direction selections available for projection.""" + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononBandHandler.selections, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PhononBandHandler.from_data, + PhononBandHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/phonon_dos.py b/src/py4vasp/_calculation/phonon_dos.py index bd2091cbc..3bb2fc45d 100644 --- a/src/py4vasp/_calculation/phonon_dos.py +++ b/src/py4vasp/_calculation/phonon_dos.py @@ -1,237 +1,237 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import pathlib - -from py4vasp import raw -from py4vasp._calculation import phonon -from py4vasp._calculation._stoichiometry import StoichiometryHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import PhononDos_DB -from py4vasp._third_party import graph -from py4vasp._util import check, documentation, index, select - - -class PhononDosHandler: - """Handler for phonon DOS data.""" - - def __init__(self, raw_phonon_dos: raw.PhononDos): - self._raw_phonon_dos = raw_phonon_dos - - @classmethod - def from_data(cls, raw_phonon_dos: raw.PhononDos) -> "PhononDosHandler": - return cls(raw_phonon_dos) - - def __str__(self) -> str: - energies = self._raw_phonon_dos.energies - stoichiometry = self._stoichiometry() - return f"""phonon DOS: - [{energies[0]:0.2f}, {energies[-1]:0.2f}] mesh with {len(energies)} points - {3 * stoichiometry.number_atoms()} modes - {stoichiometry}""" - - def to_dict(self, selection=None) -> dict: - """Read the phonon DOS into a dictionary. - - Parameters - ---------- - selection : str - A string specifying the projection of the phonon modes onto atoms and - directions. See `selections` for available options. - - Returns - ------- - dict - Contains the energies at which the phonon DOS was computed. The total - DOS is returned and any possible projected DOS selected by the *selection* - argument. - """ - return { - "energies": self._raw_phonon_dos.energies[:], - "total": self._raw_phonon_dos.dos[:], - **self._read_data(selection), - } - - def to_database(self) -> dict: - energy_min = ( - float(self._raw_phonon_dos.energies[0]) - if not check.is_none(self._raw_phonon_dos.energies) - else None - ) - energy_max = ( - float(self._raw_phonon_dos.energies[-1]) - if not check.is_none(self._raw_phonon_dos.energies) - else None - ) - return PhononDos_DB(energy_min=energy_min, energy_max=energy_max) - - def to_graph(self, selection=None) -> graph.Graph: - data = self.to_dict(selection) - return graph.Graph( - series=list(_series(data)), - xlabel="ω (THz)", - ylabel="DOS (1/THz)", - ) - - def selections(self) -> dict: - atoms = self._init_atom_dict().keys() - return { - "atom": sorted(atoms, key=self._sort_key), - "direction": ["x", "y", "z"], - } - - def _stoichiometry(self) -> StoichiometryHandler: - return StoichiometryHandler.from_data(self._raw_phonon_dos.stoichiometry) - - def _read_data(self, selection) -> dict: - if not selection: - return {} - maps = {0: self._init_atom_dict(), 1: self._init_direction_dict()} - selector = index.Selector( - maps, self._raw_phonon_dos.projections, use_number_labels=True - ) - tree = select.Tree.from_selection(selection) - return {selector.label(sel): selector[sel] for sel in tree.selections()} - - def _init_atom_dict(self) -> dict: - return { - key: value.indices - for key, value in self._stoichiometry().read().items() - if key != select.all - } - - def _init_direction_dict(self) -> dict: - return { - "x": slice(0, 1), - "y": slice(1, 2), - "z": slice(2, 3), - } - - def _sort_key(self, key) -> bool: - return key.isdecimal() - - -@quantity("dos", group="phonon") -class PhononDos(graph.Mixin): - """The phonon density of states (DOS) describes the number of modes per energy. - - The phonon density of states (DOS) is a representation of the distribution of - phonons in a material across different frequencies. It provides a histogram of the - number of phonon states per frequency interval. Peaks and features in the DOS reveal - the density of vibrational modes at specific frequencies. One can related these - properties to study e.g. the heat capacity or the thermal conductivity. - - Projecting the phonon density of states (DOS) onto specific atoms highlights the - contribution of each atomic species to the vibrational spectrum. This analysis helps - to understand the role of individual elements' impact on the material's thermal and - mechanical properties. The projected phonon DOS can guide towards engineering these - properties by substitution of specific atoms. Additionally, the atom-specific - projection allows for the identification of localized modes or vibrations associated - with specific atomic species. - """ - - def __init__(self, source, quantity_name: str = "phonon_dos"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_phonon_dos: raw.PhononDos) -> "PhononDos": - return cls(source=DataSource(raw_phonon_dos)) - - def _handler_factory(self, raw): - return PhononDosHandler.from_data(raw) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononDosHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - @documentation.format(selection=phonon.selection_doc) - def read(self, selection: str | None = None) -> dict: - """Read the phonon DOS into a dictionary. - - Parameters - ---------- - {selection} - - Returns - ------- - dict - Contains the energies at which the phonon DOS was computed. The total - DOS is returned and any possible projected DOS selected by the *selection* - argument. - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononDosHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - @documentation.format(selection=phonon.selection_doc) - def to_graph(self, selection: str | None = None) -> graph.Graph: - """Generate a graph of the selected phonon DOS. - - Parameters - ---------- - {selection} - - Returns - ------- - Graph - The graph contains the total DOS. If a selection is given, in addition the - projected DOS is shown. - """ - return merge_graphs( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononDosHandler.to_graph, - ) - - def selections(self, selection=None) -> dict: - """Return atom and direction selections available for projection.""" - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononDosHandler.selections, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PhononDosHandler.from_data, - PhononDosHandler.to_database, - ) - - -def _series(data): - energies = data["energies"] - for name, dos in data.items(): - if name == "energies": - continue - yield graph.Series(energies, dos, name) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import pathlib + +from py4vasp import raw +from py4vasp._calculation import phonon +from py4vasp._calculation._stoichiometry import StoichiometryHandler +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import PhononDos_DB +from py4vasp._third_party import graph +from py4vasp._util import check, documentation, index, select + + +class PhononDosHandler: + """Handler for phonon DOS data.""" + + def __init__(self, raw_phonon_dos: raw.PhononDos): + self._raw_phonon_dos = raw_phonon_dos + + @classmethod + def from_data(cls, raw_phonon_dos: raw.PhononDos) -> "PhononDosHandler": + return cls(raw_phonon_dos) + + def __str__(self) -> str: + energies = self._raw_phonon_dos.energies + stoichiometry = self._stoichiometry() + return f"""phonon DOS: + [{energies[0]:0.2f}, {energies[-1]:0.2f}] mesh with {len(energies)} points + {3 * stoichiometry.number_atoms()} modes + {stoichiometry}""" + + def to_dict(self, selection=None) -> dict: + """Read the phonon DOS into a dictionary. + + Parameters + ---------- + selection : str + A string specifying the projection of the phonon modes onto atoms and + directions. See `selections` for available options. + + Returns + ------- + dict + Contains the energies at which the phonon DOS was computed. The total + DOS is returned and any possible projected DOS selected by the *selection* + argument. + """ + return { + "energies": self._raw_phonon_dos.energies[:], + "total": self._raw_phonon_dos.dos[:], + **self._read_data(selection), + } + + def to_database(self) -> dict: + energy_min = ( + float(self._raw_phonon_dos.energies[0]) + if not check.is_none(self._raw_phonon_dos.energies) + else None + ) + energy_max = ( + float(self._raw_phonon_dos.energies[-1]) + if not check.is_none(self._raw_phonon_dos.energies) + else None + ) + return PhononDos_DB(energy_min=energy_min, energy_max=energy_max) + + def to_graph(self, selection=None) -> graph.Graph: + data = self.to_dict(selection) + return graph.Graph( + series=list(_series(data)), + xlabel="ω (THz)", + ylabel="DOS (1/THz)", + ) + + def selections(self) -> dict: + atoms = self._init_atom_dict().keys() + return { + "atom": sorted(atoms, key=self._sort_key), + "direction": ["x", "y", "z"], + } + + def _stoichiometry(self) -> StoichiometryHandler: + return StoichiometryHandler.from_data(self._raw_phonon_dos.stoichiometry) + + def _read_data(self, selection) -> dict: + if not selection: + return {} + maps = {0: self._init_atom_dict(), 1: self._init_direction_dict()} + selector = index.Selector( + maps, self._raw_phonon_dos.projections, use_number_labels=True + ) + tree = select.Tree.from_selection(selection) + return {selector.label(sel): selector[sel] for sel in tree.selections()} + + def _init_atom_dict(self) -> dict: + return { + key: value.indices + for key, value in self._stoichiometry().read().items() + if key != select.all + } + + def _init_direction_dict(self) -> dict: + return { + "x": slice(0, 1), + "y": slice(1, 2), + "z": slice(2, 3), + } + + def _sort_key(self, key) -> bool: + return key.isdecimal() + + +@quantity("dos", group="phonon") +class PhononDos(graph.Mixin): + """The phonon density of states (DOS) describes the number of modes per energy. + + The phonon density of states (DOS) is a representation of the distribution of + phonons in a material across different frequencies. It provides a histogram of the + number of phonon states per frequency interval. Peaks and features in the DOS reveal + the density of vibrational modes at specific frequencies. One can related these + properties to study e.g. the heat capacity or the thermal conductivity. + + Projecting the phonon density of states (DOS) onto specific atoms highlights the + contribution of each atomic species to the vibrational spectrum. This analysis helps + to understand the role of individual elements' impact on the material's thermal and + mechanical properties. The projected phonon DOS can guide towards engineering these + properties by substitution of specific atoms. Additionally, the atom-specific + projection allows for the identification of localized modes or vibrations associated + with specific atomic species. + """ + + def __init__(self, source, quantity_name: str = "phonon_dos"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_phonon_dos: raw.PhononDos) -> "PhononDos": + return cls(source=DataSource(raw_phonon_dos)) + + def _handler_factory(self, raw): + return PhononDosHandler.from_data(raw) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononDosHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + @documentation.format(selection=phonon.selection_doc) + def read(self, selection: str | None = None) -> dict: + """Read the phonon DOS into a dictionary. + + Parameters + ---------- + {selection} + + Returns + ------- + dict + Contains the energies at which the phonon DOS was computed. The total + DOS is returned and any possible projected DOS selected by the *selection* + argument. + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononDosHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + @documentation.format(selection=phonon.selection_doc) + def to_graph(self, selection: str | None = None) -> graph.Graph: + """Generate a graph of the selected phonon DOS. + + Parameters + ---------- + {selection} + + Returns + ------- + Graph + The graph contains the total DOS. If a selection is given, in addition the + projected DOS is shown. + """ + return merge_graphs( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononDosHandler.to_graph, + ) + + def selections(self, selection=None) -> dict: + """Return atom and direction selections available for projection.""" + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononDosHandler.selections, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PhononDosHandler.from_data, + PhononDosHandler.to_database, + ) + + +def _series(data): + energies = data["energies"] + for name, dos in data.items(): + if name == "energies": + continue + yield graph.Series(energies, dos, name) diff --git a/src/py4vasp/_calculation/phonon_mode.py b/src/py4vasp/_calculation/phonon_mode.py index b16daff4e..19e1a2561 100644 --- a/src/py4vasp/_calculation/phonon_mode.py +++ b/src/py4vasp/_calculation/phonon_mode.py @@ -1,160 +1,160 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import numpy as np - -from py4vasp import raw -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw.data_db import PhononMode_DB -from py4vasp._util import check, convert - - -class PhononModeHandler: - """Handler for phonon mode data — performs all data access and transformation logic.""" - - def __init__(self, raw_phonon_mode: raw.PhononMode): - self._raw_phonon_mode = raw_phonon_mode - - @classmethod - def from_data(cls, raw_phonon_mode: raw.PhononMode) -> "PhononModeHandler": - return cls(raw_phonon_mode) - - def __str__(self) -> str: - phonon_frequencies = "\n".join( - self._frequency_to_string(index, frequency) - for index, frequency in enumerate(self.frequencies()) - ) - return f"""\ - Eigenvalues of the dynamical matrix - ----------------------------------- -{phonon_frequencies} -""" - - def to_dict(self) -> dict: - return { - "structure": self._structure().to_dict(), - "frequencies": self.frequencies(), - "eigenvectors": self._raw_phonon_mode.eigenvectors[:], - } - - def to_database(self) -> dict: - frequencies = ( - self.frequencies() - if not check.is_none(self._raw_phonon_mode.frequencies) - else None - ) - frequencies_real_max = ( - float(np.max(frequencies.real)) if frequencies is not None else None - ) - frequencies_imag_max = ( - float(np.max(frequencies.imag)) if frequencies is not None else None - ) - return PhononMode_DB( - frequencies_real_max=frequencies_real_max, - frequencies_imag_max=frequencies_imag_max, - ) - - def frequencies(self) -> np.ndarray: - """Read the phonon frequencies as a numpy array.""" - return convert.to_complex(self._raw_phonon_mode.frequencies[:]) - - def _structure(self) -> StructureHandler: - return StructureHandler.from_data(self._raw_phonon_mode.structure) - - def _frequency_to_string(self, index, frequency) -> str: - if frequency.real >= frequency.imag: - label = f"{index + 1:4} f " - else: - label = f"{index + 1:4} f/i" - frequency = np.abs(frequency) - freq_meV = f"{frequency * 1000:12.6f} meV" - eV_to_THz = 241.798934781 - freq_THz = f"{frequency * eV_to_THz:11.6f} THz" - freq_2PiTHz = f"{2 * np.pi * frequency * eV_to_THz:12.6f} 2PiTHz" - eV_to_cm1 = 8065.610420 - freq_cm1 = f"{frequency * eV_to_cm1:12.6f} cm-1" - return f"{label}= {freq_THz} {freq_2PiTHz}{freq_cm1} {freq_meV}" - - -@quantity("mode", group="phonon") -class PhononMode: - """Describes a collective vibration of atoms in a crystal. - - A phonon mode represents a specific way in which atoms in a solid oscillate - around their equilibrium positions. Each mode is characterized by a frequency - and a displacement pattern that shows how atoms move relative to each other. - Low-frequency modes correspond to long-wavelength vibrations, while - high-frequency modes involve more localized atomic motion.""" - - def __init__(self, source, quantity_name: str = "phonon_mode"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_phonon_mode: raw.PhononMode) -> "PhononMode": - return cls(source=DataSource(raw_phonon_mode)) - - def _handler_factory(self, raw): - return PhononModeHandler.from_data(raw) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PhononModeHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Read structure data and properties of the phonon mode into a dictionary. - - The frequency and eigenvector describe with how atoms move under the influence - of a particular phonon mode. Structural information is added to understand - what the displacement correspond to. - - Returns - ------- - dict - Structural information, phonon frequencies and eigenvectors. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - PhononModeHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - def frequencies(self) -> np.ndarray: - """Read the phonon frequencies as a numpy array.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - PhononModeHandler.frequencies, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PhononModeHandler.from_data, - PhononModeHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import numpy as np + +from py4vasp import raw +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw.data_db import PhononMode_DB +from py4vasp._util import check, convert + + +class PhononModeHandler: + """Handler for phonon mode data — performs all data access and transformation logic.""" + + def __init__(self, raw_phonon_mode: raw.PhononMode): + self._raw_phonon_mode = raw_phonon_mode + + @classmethod + def from_data(cls, raw_phonon_mode: raw.PhononMode) -> "PhononModeHandler": + return cls(raw_phonon_mode) + + def __str__(self) -> str: + phonon_frequencies = "\n".join( + self._frequency_to_string(index, frequency) + for index, frequency in enumerate(self.frequencies()) + ) + return f"""\ + Eigenvalues of the dynamical matrix + ----------------------------------- +{phonon_frequencies} +""" + + def to_dict(self) -> dict: + return { + "structure": self._structure().to_dict(), + "frequencies": self.frequencies(), + "eigenvectors": self._raw_phonon_mode.eigenvectors[:], + } + + def to_database(self) -> dict: + frequencies = ( + self.frequencies() + if not check.is_none(self._raw_phonon_mode.frequencies) + else None + ) + frequencies_real_max = ( + float(np.max(frequencies.real)) if frequencies is not None else None + ) + frequencies_imag_max = ( + float(np.max(frequencies.imag)) if frequencies is not None else None + ) + return PhononMode_DB( + frequencies_real_max=frequencies_real_max, + frequencies_imag_max=frequencies_imag_max, + ) + + def frequencies(self) -> np.ndarray: + """Read the phonon frequencies as a numpy array.""" + return convert.to_complex(self._raw_phonon_mode.frequencies[:]) + + def _structure(self) -> StructureHandler: + return StructureHandler.from_data(self._raw_phonon_mode.structure) + + def _frequency_to_string(self, index, frequency) -> str: + if frequency.real >= frequency.imag: + label = f"{index + 1:4} f " + else: + label = f"{index + 1:4} f/i" + frequency = np.abs(frequency) + freq_meV = f"{frequency * 1000:12.6f} meV" + eV_to_THz = 241.798934781 + freq_THz = f"{frequency * eV_to_THz:11.6f} THz" + freq_2PiTHz = f"{2 * np.pi * frequency * eV_to_THz:12.6f} 2PiTHz" + eV_to_cm1 = 8065.610420 + freq_cm1 = f"{frequency * eV_to_cm1:12.6f} cm-1" + return f"{label}= {freq_THz} {freq_2PiTHz}{freq_cm1} {freq_meV}" + + +@quantity("mode", group="phonon") +class PhononMode: + """Describes a collective vibration of atoms in a crystal. + + A phonon mode represents a specific way in which atoms in a solid oscillate + around their equilibrium positions. Each mode is characterized by a frequency + and a displacement pattern that shows how atoms move relative to each other. + Low-frequency modes correspond to long-wavelength vibrations, while + high-frequency modes involve more localized atomic motion.""" + + def __init__(self, source, quantity_name: str = "phonon_mode"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_phonon_mode: raw.PhononMode) -> "PhononMode": + return cls(source=DataSource(raw_phonon_mode)) + + def _handler_factory(self, raw): + return PhononModeHandler.from_data(raw) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PhononModeHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Read structure data and properties of the phonon mode into a dictionary. + + The frequency and eigenvector describe with how atoms move under the influence + of a particular phonon mode. Structural information is added to understand + what the displacement correspond to. + + Returns + ------- + dict + Structural information, phonon frequencies and eigenvectors. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + PhononModeHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + def frequencies(self) -> np.ndarray: + """Read the phonon frequencies as a numpy array.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + PhononModeHandler.frequencies, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PhononModeHandler.from_data, + PhononModeHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/piezoelectric_tensor.py b/src/py4vasp/_calculation/piezoelectric_tensor.py index d7b931aa4..d8e23c5ef 100644 --- a/src/py4vasp/_calculation/piezoelectric_tensor.py +++ b/src/py4vasp/_calculation/piezoelectric_tensor.py @@ -1,398 +1,398 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from contextlib import suppress - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation.cell import CellHandler -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import PiezoelectricTensor_DB -from py4vasp._util import check -from py4vasp._util.tensor import symmetry_reduce, tensor_constants - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - AttributeError, - TypeError, - ValueError, - IndexError, -) - - -class PiezoelectricTensorHandler: - """Handler for the piezoelectric_tensor quantity. Works with exactly one raw.PiezoelectricTensor object.""" - - def __init__(self, raw_piezoelectric_tensor: raw.PiezoelectricTensor): - self._raw_piezoelectric_tensor = raw_piezoelectric_tensor - - @classmethod - def from_data( - cls, raw_piezoelectric_tensor: raw.PiezoelectricTensor - ) -> "PiezoelectricTensorHandler": - return cls(raw_piezoelectric_tensor) - - def __str__(self): - data = self.to_dict() - return f"""Piezoelectric tensor (C/m²) - XX YY ZZ XY YZ ZX ---------------------------------------------------------------------------- -{_tensor_to_string(data["clamped_ion"], "clamped-ion")} -{_tensor_to_string(data["relaxed_ion"], "relaxed-ion")}""" - - def to_dict(self) -> dict: - """Read the ionic and electronic contribution to the piezoelectric tensor - into a dictionary. - - It will combine both terms as the total piezoelectric tensor (relaxed_ion) - but also give the pure electronic contribution, so that you can separate the - parts. - - Returns - ------- - dict - The clamped ion and relaxed ion data for the piezoelectric tensor. - """ - electron_data = self._raw_piezoelectric_tensor.electron[:] - return { - "clamped_ion": electron_data, - "relaxed_ion": electron_data + self._raw_piezoelectric_tensor.ion[:], - } - - def to_database(self) -> dict: - reduced_tensor_x, reduced_tensor_y, reduced_tensor_z, tensor_2d = ( - [None, None, None] for _ in range(4) - ) - in_plane = [[False, False, False] for _ in range(3)] - e11, e22, e33, e_avg_abs, e_rms, e_frobenius = ( - [None, None, None] for _ in range(6) - ) - - total_tensor, electronic_tensor, ionic_tensor = None, None, None - if not check.is_none(self._raw_piezoelectric_tensor.ion) and not check.is_none( - self._raw_piezoelectric_tensor.electron - ): - total_tensor = ( - self._raw_piezoelectric_tensor.electron[:] - + self._raw_piezoelectric_tensor.ion[:] - ) - if not check.is_none(self._raw_piezoelectric_tensor.electron): - electronic_tensor = self._raw_piezoelectric_tensor.electron[:] - if not check.is_none(self._raw_piezoelectric_tensor.ion): - ionic_tensor = self._raw_piezoelectric_tensor.ion[:] - - for idt, tensor in enumerate([total_tensor, ionic_tensor, electronic_tensor]): - e_tensor = None - # Piezoelectric stress tensor e_ij (C/m^2) - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - e_tensor = _extract_tensor( - tensor - ) # 3x6 tensor, column order: XX YY ZZ YZ ZX XY - # write in default VASP order (rows: x,y,z; columns: XX,YY,ZZ,XY,YZ,ZX) - reduced_tensor_x[idt] = e_tensor[0, (0, 1, 2, 5, 3, 4)].tolist() - reduced_tensor_y[idt] = e_tensor[1, (0, 1, 2, 5, 3, 4)].tolist() - reduced_tensor_z[idt] = e_tensor[2, (0, 1, 2, 5, 3, 4)].tolist() - - if not check.is_none(self._raw_piezoelectric_tensor.cell): - cCell = CellHandler.from_data( - self._raw_piezoelectric_tensor.cell, steps=-1 - ) - in_plane[idt], lvac = _compute_2d_plane_and_conversion_factor(cCell) - if in_plane[idt] is not None and lvac is not None: - tensor_2d[idt] = e_tensor * lvac - - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - ( - e11[idt], - e22[idt], - e33[idt], - e_avg_abs[idt], - e_rms[idt], - e_frobenius[idt], - ) = _compute_bulk_quantities(e_tensor) - - return PiezoelectricTensor_DB( - total_3d_tensor_x=reduced_tensor_x[0], - total_3d_tensor_y=reduced_tensor_y[0], - total_3d_tensor_z=reduced_tensor_z[0], - total_3d_piezoelectric_stress_coefficient_x=e11[0], - total_3d_piezoelectric_stress_coefficient_y=e22[0], - total_3d_piezoelectric_stress_coefficient_z=e33[0], - total_3d_mean_absolute=e_avg_abs[0], - total_3d_rms=e_rms[0], - total_3d_frobenius_norm=e_frobenius[0], - total_2d_tensor_x=( - ( - tensor_2d[0][0] - if (in_plane[0] is not None and in_plane[0][0]) - else None - ) - if tensor_2d[0] is not None - else None - ), - total_2d_tensor_y=( - ( - tensor_2d[0][1] - if (in_plane[0] is not None and in_plane[0][1]) - else None - ) - if tensor_2d[0] is not None - else None - ), - total_2d_tensor_z=( - ( - tensor_2d[0][2] - if (in_plane[0] is not None and in_plane[0][2]) - else None - ) - if tensor_2d[0] is not None - else None - ), - ionic_3d_tensor_x=reduced_tensor_x[1], - ionic_3d_tensor_y=reduced_tensor_y[1], - ionic_3d_tensor_z=reduced_tensor_z[1], - ionic_3d_piezoelectric_stress_coefficient_x=e11[1], - ionic_3d_piezoelectric_stress_coefficient_y=e22[1], - ionic_3d_piezoelectric_stress_coefficient_z=e33[1], - ionic_3d_mean_absolute=e_avg_abs[1], - ionic_3d_rms=e_rms[1], - ionic_3d_frobenius_norm=e_frobenius[1], - ionic_2d_tensor_x=( - ( - tensor_2d[1][0] - if (in_plane[1] is not None and in_plane[1][0]) - else None - ) - if tensor_2d[1] is not None - else None - ), - ionic_2d_tensor_y=( - ( - tensor_2d[1][1] - if (in_plane[1] is not None and in_plane[1][1]) - else None - ) - if tensor_2d[1] is not None - else None - ), - ionic_2d_tensor_z=( - ( - tensor_2d[1][2] - if (in_plane[1] is not None and in_plane[1][2]) - else None - ) - if tensor_2d[1] is not None - else None - ), - electronic_3d_tensor_x=reduced_tensor_x[2], - electronic_3d_tensor_y=reduced_tensor_y[2], - electronic_3d_tensor_z=reduced_tensor_z[2], - electronic_3d_piezoelectric_stress_coefficient_x=e11[2], - electronic_3d_piezoelectric_stress_coefficient_y=e22[2], - electronic_3d_piezoelectric_stress_coefficient_z=e33[2], - electronic_3d_mean_absolute=e_avg_abs[2], - electronic_3d_rms=e_rms[2], - electronic_3d_frobenius_norm=e_frobenius[2], - electronic_2d_tensor_x=( - ( - tensor_2d[2][0] - if (in_plane[2] is not None and in_plane[2][0]) - else None - ) - if tensor_2d[2] is not None - else None - ), - electronic_2d_tensor_y=( - ( - tensor_2d[2][1] - if (in_plane[2] is not None and in_plane[2][1]) - else None - ) - if tensor_2d[2] is not None - else None - ), - electronic_2d_tensor_z=( - ( - tensor_2d[2][2] - if (in_plane[2] is not None and in_plane[2][2]) - else None - ) - if tensor_2d[2] is not None - else None - ), - ) - - -@quantity("piezoelectric_tensor") -class PiezoelectricTensor: - """The piezoelectric tensor is the derivative of the energy with respect to strain and field. - - The piezoelectric tensor represents the coupling between mechanical stress and - electrical polarization in a material. VASP computes the piezoelectric tensor with - a linear response calculation. The piezoelectric tensor is a 3x3 matrix that relates - the three components of stress to the three components of polarization. - Specifically, it describes how the application of mechanical stress induces an - electric polarization and, conversely, how an applied electric field results in - a deformation. - - The piezoelectric tensor helps to characterize the efficiency and anisotropy of the - piezoelectric response. A large piezoelectric tensor is useful e.g. for sensors - and actuators. Moreover, the tensor's symmetry properties are coupled to the crystal - structure and symmetry. Therefore a mismatch of the symmetry properties between - calculations and experiment can reveal underlying flaws in the characterization of - the crystal structure. - """ - - def __init__(self, source, quantity_name: str = "piezoelectric_tensor"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data( - cls, raw_piezoelectric_tensor: raw.PiezoelectricTensor - ) -> "PiezoelectricTensor": - """Create a PiezoelectricTensor dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_piezoelectric_tensor)) - - def read(self) -> dict: - """Read the ionic and electronic contribution to the piezoelectric tensor - into a dictionary. - - It will combine both terms as the total piezoelectric tensor (relaxed_ion) - but also give the pure electronic contribution, so that you can separate the - parts. - - Returns - ------- - dict - The clamped ion and relaxed ion data for the piezoelectric tensor. - """ - return merge_default( - self._source, - self._quantity_name, - None, - PiezoelectricTensorHandler.from_data, - PiezoelectricTensorHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - PiezoelectricTensorHandler.from_data, - PiezoelectricTensorHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PiezoelectricTensorHandler.from_data, - PiezoelectricTensorHandler.to_database, - ) - - -def _tensor_to_string(tensor, label): - compact_tensor = symmetry_reduce(tensor.T).T - line = lambda dir_, vec: dir_ + " " + " ".join(f"{x:11.5f}" for x in vec) - directions = (" x", " y", " z") - lines = (line(dir_, vec) for dir_, vec in zip(directions, compact_tensor)) - return f"{label:^75}".rstrip() + "\n" + "\n".join(lines) - - -def _extract_tensor(raw_tensor): - voigt_indices = { - (0, 0): 0, # XX - (1, 1): 1, # YY - (2, 2): 2, # ZZ - (1, 2): 3, # YZ - (0, 2): 4, # ZX - (0, 1): 5, # XY - } - C_voigt = np.zeros((3, 6)) - for i in range(3): - for j in range(i, 3): - for k in range(3): - column = voigt_indices.get((i, j)) - C_voigt[k, column] = raw_tensor[i, j, k] - - return C_voigt - - -def _compute_2d_piezoelectric(e_tensor: np.ndarray, cell_: CellHandler) -> np.ndarray: - """ - Convert 3D piezoelectric stress tensor (C/m^2) to 2D (C/m). - """ - # TODO migrate finding vacuum direction to structure - in_plane, l_vac = _compute_2d_plane_and_conversion_factor(cell_) - return e_tensor[in_plane, :] * l_vac - - -def _compute_2d_plane_and_conversion_factor( - cell_: CellHandler, -) -> tuple[list[bool], float]: - """ - Identify the 2D plane (in-plane directions) and compute the conversion factor - from 3D piezoelectric tensor (C/m^2) to 2D (C/m). - """ - vac_dir = cell_._find_likely_vacuum_direction() - if vac_dir is None: - return None, None - in_plane = [i != vac_dir for i in range(3)] - l_vac = np.linalg.norm(cell_.lattice_vectors()[vac_dir]) * 1e-10 # Å → m - return in_plane, l_vac - - -def _compute_bulk_quantities( - e_tensor: np.ndarray, -) -> tuple[float, float, float, float, float, float]: - """ - Scalar bulk measures from a VASP piezoelectric stress tensor (3x6). - Units: C/m^2 - - e11: Piezoelectric stress coefficient: polarization along x induced by normal strain along x (∂Px/∂εxx). - - e22: Piezoelectric stress coefficient: polarization along y induced by normal strain along y (∂Py/∂εyy). - - e33: Piezoelectric stress coefficient: polarization along z induced by normal strain along z (∂Pz/∂εzz); - *** sign indicates direction of polarization relative to applied strain. - - e_avg_abs: - - Mean absolute value of all piezoelectric tensor components; - - a scalar descriptor of the overall piezoelectric response magnitude (not a fundamental constant). - - e_rms: - - Root-mean-square of piezoelectric tensor components; - - emphasizes larger tensor elements and provides a normalized measure of overall piezoelectric strength. - - e_frobenius: - - Frobenius norm of the piezoelectric tensor; - - rotation-invariant total magnitude of the piezoelectric response, commonly used for materials screening. - - """ - if (e_tensor is None) or (e_tensor.shape != (3, 6)): - return None, None, None, None, None, None - e11 = e_tensor[0, 0] - e22 = e_tensor[1, 1] - e33 = e_tensor[2, 2] - e_avg_abs = np.mean(np.abs(e_tensor)) - e_rms = np.sqrt(np.mean(e_tensor**2)) - e_frobenius = np.linalg.norm(e_tensor) - - return e11, e22, e33, e_avg_abs, e_rms, e_frobenius +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from contextlib import suppress + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation.cell import CellHandler +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import PiezoelectricTensor_DB +from py4vasp._util import check +from py4vasp._util.tensor import symmetry_reduce, tensor_constants + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + AttributeError, + TypeError, + ValueError, + IndexError, +) + + +class PiezoelectricTensorHandler: + """Handler for the piezoelectric_tensor quantity. Works with exactly one raw.PiezoelectricTensor object.""" + + def __init__(self, raw_piezoelectric_tensor: raw.PiezoelectricTensor): + self._raw_piezoelectric_tensor = raw_piezoelectric_tensor + + @classmethod + def from_data( + cls, raw_piezoelectric_tensor: raw.PiezoelectricTensor + ) -> "PiezoelectricTensorHandler": + return cls(raw_piezoelectric_tensor) + + def __str__(self): + data = self.to_dict() + return f"""Piezoelectric tensor (C/m²) + XX YY ZZ XY YZ ZX +--------------------------------------------------------------------------- +{_tensor_to_string(data["clamped_ion"], "clamped-ion")} +{_tensor_to_string(data["relaxed_ion"], "relaxed-ion")}""" + + def to_dict(self) -> dict: + """Read the ionic and electronic contribution to the piezoelectric tensor + into a dictionary. + + It will combine both terms as the total piezoelectric tensor (relaxed_ion) + but also give the pure electronic contribution, so that you can separate the + parts. + + Returns + ------- + dict + The clamped ion and relaxed ion data for the piezoelectric tensor. + """ + electron_data = self._raw_piezoelectric_tensor.electron[:] + return { + "clamped_ion": electron_data, + "relaxed_ion": electron_data + self._raw_piezoelectric_tensor.ion[:], + } + + def to_database(self) -> dict: + reduced_tensor_x, reduced_tensor_y, reduced_tensor_z, tensor_2d = ( + [None, None, None] for _ in range(4) + ) + in_plane = [[False, False, False] for _ in range(3)] + e11, e22, e33, e_avg_abs, e_rms, e_frobenius = ( + [None, None, None] for _ in range(6) + ) + + total_tensor, electronic_tensor, ionic_tensor = None, None, None + if not check.is_none(self._raw_piezoelectric_tensor.ion) and not check.is_none( + self._raw_piezoelectric_tensor.electron + ): + total_tensor = ( + self._raw_piezoelectric_tensor.electron[:] + + self._raw_piezoelectric_tensor.ion[:] + ) + if not check.is_none(self._raw_piezoelectric_tensor.electron): + electronic_tensor = self._raw_piezoelectric_tensor.electron[:] + if not check.is_none(self._raw_piezoelectric_tensor.ion): + ionic_tensor = self._raw_piezoelectric_tensor.ion[:] + + for idt, tensor in enumerate([total_tensor, ionic_tensor, electronic_tensor]): + e_tensor = None + # Piezoelectric stress tensor e_ij (C/m^2) + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + e_tensor = _extract_tensor( + tensor + ) # 3x6 tensor, column order: XX YY ZZ YZ ZX XY + # write in default VASP order (rows: x,y,z; columns: XX,YY,ZZ,XY,YZ,ZX) + reduced_tensor_x[idt] = e_tensor[0, (0, 1, 2, 5, 3, 4)].tolist() + reduced_tensor_y[idt] = e_tensor[1, (0, 1, 2, 5, 3, 4)].tolist() + reduced_tensor_z[idt] = e_tensor[2, (0, 1, 2, 5, 3, 4)].tolist() + + if not check.is_none(self._raw_piezoelectric_tensor.cell): + cCell = CellHandler.from_data( + self._raw_piezoelectric_tensor.cell, steps=-1 + ) + in_plane[idt], lvac = _compute_2d_plane_and_conversion_factor(cCell) + if in_plane[idt] is not None and lvac is not None: + tensor_2d[idt] = e_tensor * lvac + + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + ( + e11[idt], + e22[idt], + e33[idt], + e_avg_abs[idt], + e_rms[idt], + e_frobenius[idt], + ) = _compute_bulk_quantities(e_tensor) + + return PiezoelectricTensor_DB( + total_3d_tensor_x=reduced_tensor_x[0], + total_3d_tensor_y=reduced_tensor_y[0], + total_3d_tensor_z=reduced_tensor_z[0], + total_3d_piezoelectric_stress_coefficient_x=e11[0], + total_3d_piezoelectric_stress_coefficient_y=e22[0], + total_3d_piezoelectric_stress_coefficient_z=e33[0], + total_3d_mean_absolute=e_avg_abs[0], + total_3d_rms=e_rms[0], + total_3d_frobenius_norm=e_frobenius[0], + total_2d_tensor_x=( + ( + tensor_2d[0][0] + if (in_plane[0] is not None and in_plane[0][0]) + else None + ) + if tensor_2d[0] is not None + else None + ), + total_2d_tensor_y=( + ( + tensor_2d[0][1] + if (in_plane[0] is not None and in_plane[0][1]) + else None + ) + if tensor_2d[0] is not None + else None + ), + total_2d_tensor_z=( + ( + tensor_2d[0][2] + if (in_plane[0] is not None and in_plane[0][2]) + else None + ) + if tensor_2d[0] is not None + else None + ), + ionic_3d_tensor_x=reduced_tensor_x[1], + ionic_3d_tensor_y=reduced_tensor_y[1], + ionic_3d_tensor_z=reduced_tensor_z[1], + ionic_3d_piezoelectric_stress_coefficient_x=e11[1], + ionic_3d_piezoelectric_stress_coefficient_y=e22[1], + ionic_3d_piezoelectric_stress_coefficient_z=e33[1], + ionic_3d_mean_absolute=e_avg_abs[1], + ionic_3d_rms=e_rms[1], + ionic_3d_frobenius_norm=e_frobenius[1], + ionic_2d_tensor_x=( + ( + tensor_2d[1][0] + if (in_plane[1] is not None and in_plane[1][0]) + else None + ) + if tensor_2d[1] is not None + else None + ), + ionic_2d_tensor_y=( + ( + tensor_2d[1][1] + if (in_plane[1] is not None and in_plane[1][1]) + else None + ) + if tensor_2d[1] is not None + else None + ), + ionic_2d_tensor_z=( + ( + tensor_2d[1][2] + if (in_plane[1] is not None and in_plane[1][2]) + else None + ) + if tensor_2d[1] is not None + else None + ), + electronic_3d_tensor_x=reduced_tensor_x[2], + electronic_3d_tensor_y=reduced_tensor_y[2], + electronic_3d_tensor_z=reduced_tensor_z[2], + electronic_3d_piezoelectric_stress_coefficient_x=e11[2], + electronic_3d_piezoelectric_stress_coefficient_y=e22[2], + electronic_3d_piezoelectric_stress_coefficient_z=e33[2], + electronic_3d_mean_absolute=e_avg_abs[2], + electronic_3d_rms=e_rms[2], + electronic_3d_frobenius_norm=e_frobenius[2], + electronic_2d_tensor_x=( + ( + tensor_2d[2][0] + if (in_plane[2] is not None and in_plane[2][0]) + else None + ) + if tensor_2d[2] is not None + else None + ), + electronic_2d_tensor_y=( + ( + tensor_2d[2][1] + if (in_plane[2] is not None and in_plane[2][1]) + else None + ) + if tensor_2d[2] is not None + else None + ), + electronic_2d_tensor_z=( + ( + tensor_2d[2][2] + if (in_plane[2] is not None and in_plane[2][2]) + else None + ) + if tensor_2d[2] is not None + else None + ), + ) + + +@quantity("piezoelectric_tensor") +class PiezoelectricTensor: + """The piezoelectric tensor is the derivative of the energy with respect to strain and field. + + The piezoelectric tensor represents the coupling between mechanical stress and + electrical polarization in a material. VASP computes the piezoelectric tensor with + a linear response calculation. The piezoelectric tensor is a 3x3 matrix that relates + the three components of stress to the three components of polarization. + Specifically, it describes how the application of mechanical stress induces an + electric polarization and, conversely, how an applied electric field results in + a deformation. + + The piezoelectric tensor helps to characterize the efficiency and anisotropy of the + piezoelectric response. A large piezoelectric tensor is useful e.g. for sensors + and actuators. Moreover, the tensor's symmetry properties are coupled to the crystal + structure and symmetry. Therefore a mismatch of the symmetry properties between + calculations and experiment can reveal underlying flaws in the characterization of + the crystal structure. + """ + + def __init__(self, source, quantity_name: str = "piezoelectric_tensor"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data( + cls, raw_piezoelectric_tensor: raw.PiezoelectricTensor + ) -> "PiezoelectricTensor": + """Create a PiezoelectricTensor dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_piezoelectric_tensor)) + + def read(self) -> dict: + """Read the ionic and electronic contribution to the piezoelectric tensor + into a dictionary. + + It will combine both terms as the total piezoelectric tensor (relaxed_ion) + but also give the pure electronic contribution, so that you can separate the + parts. + + Returns + ------- + dict + The clamped ion and relaxed ion data for the piezoelectric tensor. + """ + return merge_default( + self._source, + self._quantity_name, + None, + PiezoelectricTensorHandler.from_data, + PiezoelectricTensorHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + PiezoelectricTensorHandler.from_data, + PiezoelectricTensorHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PiezoelectricTensorHandler.from_data, + PiezoelectricTensorHandler.to_database, + ) + + +def _tensor_to_string(tensor, label): + compact_tensor = symmetry_reduce(tensor.T).T + line = lambda dir_, vec: dir_ + " " + " ".join(f"{x:11.5f}" for x in vec) + directions = (" x", " y", " z") + lines = (line(dir_, vec) for dir_, vec in zip(directions, compact_tensor)) + return f"{label:^75}".rstrip() + "\n" + "\n".join(lines) + + +def _extract_tensor(raw_tensor): + voigt_indices = { + (0, 0): 0, # XX + (1, 1): 1, # YY + (2, 2): 2, # ZZ + (1, 2): 3, # YZ + (0, 2): 4, # ZX + (0, 1): 5, # XY + } + C_voigt = np.zeros((3, 6)) + for i in range(3): + for j in range(i, 3): + for k in range(3): + column = voigt_indices.get((i, j)) + C_voigt[k, column] = raw_tensor[i, j, k] + + return C_voigt + + +def _compute_2d_piezoelectric(e_tensor: np.ndarray, cell_: CellHandler) -> np.ndarray: + """ + Convert 3D piezoelectric stress tensor (C/m^2) to 2D (C/m). + """ + # TODO migrate finding vacuum direction to structure + in_plane, l_vac = _compute_2d_plane_and_conversion_factor(cell_) + return e_tensor[in_plane, :] * l_vac + + +def _compute_2d_plane_and_conversion_factor( + cell_: CellHandler, +) -> tuple[list[bool], float]: + """ + Identify the 2D plane (in-plane directions) and compute the conversion factor + from 3D piezoelectric tensor (C/m^2) to 2D (C/m). + """ + vac_dir = cell_._find_likely_vacuum_direction() + if vac_dir is None: + return None, None + in_plane = [i != vac_dir for i in range(3)] + l_vac = np.linalg.norm(cell_.lattice_vectors()[vac_dir]) * 1e-10 # Å → m + return in_plane, l_vac + + +def _compute_bulk_quantities( + e_tensor: np.ndarray, +) -> tuple[float, float, float, float, float, float]: + """ + Scalar bulk measures from a VASP piezoelectric stress tensor (3x6). + Units: C/m^2 + + e11: Piezoelectric stress coefficient: polarization along x induced by normal strain along x (∂Px/∂εxx). + + e22: Piezoelectric stress coefficient: polarization along y induced by normal strain along y (∂Py/∂εyy). + + e33: Piezoelectric stress coefficient: polarization along z induced by normal strain along z (∂Pz/∂εzz); + *** sign indicates direction of polarization relative to applied strain. + + e_avg_abs: + - Mean absolute value of all piezoelectric tensor components; + - a scalar descriptor of the overall piezoelectric response magnitude (not a fundamental constant). + + e_rms: + - Root-mean-square of piezoelectric tensor components; + - emphasizes larger tensor elements and provides a normalized measure of overall piezoelectric strength. + + e_frobenius: + - Frobenius norm of the piezoelectric tensor; + - rotation-invariant total magnitude of the piezoelectric response, commonly used for materials screening. + + """ + if (e_tensor is None) or (e_tensor.shape != (3, 6)): + return None, None, None, None, None, None + e11 = e_tensor[0, 0] + e22 = e_tensor[1, 1] + e33 = e_tensor[2, 2] + e_avg_abs = np.mean(np.abs(e_tensor)) + e_rms = np.sqrt(np.mean(e_tensor**2)) + e_frobenius = np.linalg.norm(e_tensor) + + return e11, e22, e33, e_avg_abs, e_rms, e_frobenius diff --git a/src/py4vasp/_calculation/polarization.py b/src/py4vasp/_calculation/polarization.py index ca52b562d..23d8b8518 100644 --- a/src/py4vasp/_calculation/polarization.py +++ b/src/py4vasp/_calculation/polarization.py @@ -1,145 +1,145 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from contextlib import suppress - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import Polarization_DB - - -class PolarizationHandler: - """Handler for the polarization quantity. Works with exactly one raw.Polarization object.""" - - def __init__(self, raw_polarization: raw.Polarization): - self._raw_polarization = raw_polarization - - @classmethod - def from_data(cls, raw_polarization: raw.Polarization) -> "PolarizationHandler": - return cls(raw_polarization) - - def to_dict(self) -> dict: - """Read electronic and ionic polarization into a dictionary. - - Returns - ------- - dict - Contains the electronic and ionic dipole moments. - """ - return { - "electron_dipole": self._raw_polarization.electron[:], - "ion_dipole": self._raw_polarization.ion[:], - } - - def __str__(self): - vec_to_string = lambda vec: " ".join(f"{x:11.5f}" for x in vec) - return f"""Polarization (|e|Å) -------------------------------------------------------------- -ionic dipole moment: {vec_to_string(self._raw_polarization.ion[:])} -electronic dipole moment: {vec_to_string(self._raw_polarization.electron[:])}""".strip() - - def to_database(self) -> dict: - ionic_norm = None - electronic_norm = None - total_norm = None - - electron_dipole = None - ion_dipole = None - total_dipole = None - - with suppress(exception.NoData): - electron_dipole = list(self._raw_polarization.electron[:]) - ion_dipole = list(self._raw_polarization.ion[:]) - total_dipole = list( - self._raw_polarization.electron[:] + self._raw_polarization.ion[:] - ) - - ionic_norm = np.linalg.norm(self._raw_polarization.ion[:]) - electronic_norm = np.linalg.norm(self._raw_polarization.electron[:]) - total_norm = np.linalg.norm( - self._raw_polarization.electron[:] + self._raw_polarization.ion[:] - ) - - return Polarization_DB( - total_dipole_norm=total_norm, - total_dipole_moment=total_dipole, - ionic_dipole_norm=ionic_norm, - ionic_dipole_moment=ion_dipole, - electronic_dipole_norm=electronic_norm, - electronic_dipole_moment=electron_dipole, - ) - - -@quantity("polarization") -class Polarization: - """The static polarization describes the electric dipole moment per unit volume. - - Static polarization arises in a material in response to a constant external electric - field. In VASP, we compute the linear response of the system when applying a - :tag:`EFIELD`. Static polarization is a key characteristic of ferroelectric - materials that exhibit a spontaneous electric polarization that persists even in - the absence of an external electric field. - - Note that the polarization is only well defined relative to a reference - system. The absolute value can change by a polarization quantum if some - charge or ion leaves one side of the unit cell and reenters at the opposite - side. Therefore you always need to compare changes of polarization. - """ - - def __init__(self, source, quantity_name: str = "polarization"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_polarization: raw.Polarization) -> "Polarization": - """Create a Polarization dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_polarization)) - - def read(self) -> dict: - """Read electronic and ionic polarization into a dictionary. - - Returns - ------- - dict - Contains the electronic and ionic dipole moments. - """ - return merge_default( - self._source, - self._quantity_name, - None, - PolarizationHandler.from_data, - PolarizationHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - PolarizationHandler.from_data, - PolarizationHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PolarizationHandler.from_data, - PolarizationHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from contextlib import suppress + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import Polarization_DB + + +class PolarizationHandler: + """Handler for the polarization quantity. Works with exactly one raw.Polarization object.""" + + def __init__(self, raw_polarization: raw.Polarization): + self._raw_polarization = raw_polarization + + @classmethod + def from_data(cls, raw_polarization: raw.Polarization) -> "PolarizationHandler": + return cls(raw_polarization) + + def to_dict(self) -> dict: + """Read electronic and ionic polarization into a dictionary. + + Returns + ------- + dict + Contains the electronic and ionic dipole moments. + """ + return { + "electron_dipole": self._raw_polarization.electron[:], + "ion_dipole": self._raw_polarization.ion[:], + } + + def __str__(self): + vec_to_string = lambda vec: " ".join(f"{x:11.5f}" for x in vec) + return f"""Polarization (|e|Å) +------------------------------------------------------------- +ionic dipole moment: {vec_to_string(self._raw_polarization.ion[:])} +electronic dipole moment: {vec_to_string(self._raw_polarization.electron[:])}""".strip() + + def to_database(self) -> dict: + ionic_norm = None + electronic_norm = None + total_norm = None + + electron_dipole = None + ion_dipole = None + total_dipole = None + + with suppress(exception.NoData): + electron_dipole = list(self._raw_polarization.electron[:]) + ion_dipole = list(self._raw_polarization.ion[:]) + total_dipole = list( + self._raw_polarization.electron[:] + self._raw_polarization.ion[:] + ) + + ionic_norm = np.linalg.norm(self._raw_polarization.ion[:]) + electronic_norm = np.linalg.norm(self._raw_polarization.electron[:]) + total_norm = np.linalg.norm( + self._raw_polarization.electron[:] + self._raw_polarization.ion[:] + ) + + return Polarization_DB( + total_dipole_norm=total_norm, + total_dipole_moment=total_dipole, + ionic_dipole_norm=ionic_norm, + ionic_dipole_moment=ion_dipole, + electronic_dipole_norm=electronic_norm, + electronic_dipole_moment=electron_dipole, + ) + + +@quantity("polarization") +class Polarization: + """The static polarization describes the electric dipole moment per unit volume. + + Static polarization arises in a material in response to a constant external electric + field. In VASP, we compute the linear response of the system when applying a + :tag:`EFIELD`. Static polarization is a key characteristic of ferroelectric + materials that exhibit a spontaneous electric polarization that persists even in + the absence of an external electric field. + + Note that the polarization is only well defined relative to a reference + system. The absolute value can change by a polarization quantum if some + charge or ion leaves one side of the unit cell and reenters at the opposite + side. Therefore you always need to compare changes of polarization. + """ + + def __init__(self, source, quantity_name: str = "polarization"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_polarization: raw.Polarization) -> "Polarization": + """Create a Polarization dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_polarization)) + + def read(self) -> dict: + """Read electronic and ionic polarization into a dictionary. + + Returns + ------- + dict + Contains the electronic and ionic dipole moments. + """ + return merge_default( + self._source, + self._quantity_name, + None, + PolarizationHandler.from_data, + PolarizationHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + PolarizationHandler.from_data, + PolarizationHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PolarizationHandler.from_data, + PolarizationHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/potential.py b/src/py4vasp/_calculation/potential.py index 85467a4c4..95020b2e7 100644 --- a/src/py4vasp/_calculation/potential.py +++ b/src/py4vasp/_calculation/potential.py @@ -1,563 +1,563 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy -import itertools -from typing import Optional, Union - -import numpy as np - -from py4vasp import _config, exception -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Potential_DB -from py4vasp._third_party import view -from py4vasp._util import ( - check, - density, - documentation, - index, - select, - slicing, - suggest, -) - -VALID_KINDS = ("total", "ionic", "xc", "hartree") -_COMPONENTS = { - 0: ["0", "unity", "sigma_0", "scalar"], - 1: ["1", "sigma_x", "x", "sigma_1"], - 2: ["2", "sigma_y", "y", "sigma_2"], - 3: ["3", "sigma_z", "z", "sigma_3"], -} -_COMMON_PARAMETERS = f"""\ -{slicing.PARAMETERS} -supercell : int or np.ndarray - Replicate the contour plot periodically a given number of times. If you - provide two different numbers, the resulting cell will be the two remaining - lattice vectors multiplied by the specific number. -""" - - -class PotentialHandler: - """Handler for potential data — performs all data access and transformation.""" - - def __init__(self, raw_potential: raw.Potential): - self._raw_potential = raw_potential - - @classmethod - def from_data(cls, raw_potential: raw.Potential) -> "PotentialHandler": - return cls(raw_potential) - - def __str__(self) -> str: - potential = self._raw_potential.total_potential - if _is_collinear(potential): - description = "collinear potential:" - elif _is_noncollinear(potential): - description = "noncollinear potential:" - else: - description = "nonpolarized potential:" - stoichiometry = _stoichiometry.Stoichiometry.from_data( - self._raw_potential.structure.stoichiometry - ) - structure = f"structure: {stoichiometry}" - grid = f"grid: {potential.shape[3]}, {potential.shape[2]}, {potential.shape[1]}" - available = "available: " + ", ".join( - kind for kind in VALID_KINDS if not self._get_potential(kind).is_none() - ) - return "\n ".join([description, structure, grid, available]) - - def to_dict(self) -> dict: - result = {"structure": self._structure().to_dict()} - items = [self._generate_items(kind) for kind in VALID_KINDS] - result.update(itertools.chain(*items)) - return result - - def to_database(self) -> Potential_DB: - total_potential_mean = None - total_potential_mean_up = None - total_potential_mean_down = None - total_potential_mean_magnetization = None - total_potential = self._get_potential("total") - if not check.is_none(total_potential): - total_potential = np.moveaxis(total_potential, 0, -1).T - total_potential_mean = float(np.mean(total_potential[0])) - total_potential_mean_up = ( - float(np.mean(total_potential[0] + total_potential[1])) - if _is_collinear(total_potential) - else total_potential_mean / 2.0 - ) - total_potential_mean_down = ( - float(np.mean(total_potential[0] - total_potential[1])) - if _is_collinear(total_potential) - else total_potential_mean / 2.0 - ) - total_potential_mean_magnetization = ( - float(np.mean(np.linalg.norm(total_potential[1:], axis=-1))) - if _is_noncollinear(total_potential) - else None - ) - - has_potential_dict = { - f"has_{kind}_potential": not check.is_none(self._get_potential(kind)) - for kind in VALID_KINDS - } - - return Potential_DB( - **has_potential_dict, - total_potential_mean=total_potential_mean, - total_potential_mean_up=total_potential_mean_up, - total_potential_mean_down=total_potential_mean_down, - total_potential_mean_magnetization=total_potential_mean_magnetization, - ) - - def to_view( - self, - selection: str = "total", - supercell: Optional[Union[int, np.ndarray]] = None, - **user_options, - ): - potentials = dict(self._get_potentials(selection)) - isosurface = self._create_isosurface(**user_options) - viewer = self._structure().to_view(supercell) - viewer.grid_scalars = [ - view.GridQuantity( - quantity=data[np.newaxis], - label=label, - isosurfaces=[isosurface], - ) - for label, data in potentials.items() - ] - return viewer - - def to_contour( - self, - selection: str = "total", - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ): - potentials = dict(self._get_potentials(selection)) - visualizer = density.Visualizer(self._structure()) - slice_arguments = density.SliceArguments(a, b, c, normal, supercell) - return visualizer.to_contour(potentials, slice_arguments, isolevels=False) - - def to_quiver( - self, selection="total", *, a=None, b=None, c=None, normal=None, supercell=None - ): - potentials = dict(self._get_potentials(selection, is_magnetic=True)) - visualizer = density.Visualizer(self._structure()) - slice_arguments = density.SliceArguments(a, b, c, normal, supercell) - return visualizer.to_quiver(potentials, slice_arguments) - - def _structure(self): - return StructureHandler.from_data(self._raw_potential.structure) - - def _get_potentials(self, selection, is_magnetic=False): - tree = select.Tree.from_selection(selection) - for selection in tree.selections(): - kind, component = self._determine_kind_and_component(selection) - selector = self._create_selector(kind, component, is_magnetic) - component_label = component[0] if component else "" - yield self._get_label(kind, component_label), selector[component].T - - def _determine_kind_and_component(self, selection): - for kind in VALID_KINDS: - if kind in selection: - remaining = list(selection) - remaining.remove(kind) - return kind, tuple(remaining) - return "total", selection - - def _get_label(self, kind, component): - return f"{kind} potential" + (f"({component})" if component else "") - - def _create_selector(self, kind, component, is_magnetic): - if is_magnetic: - return self._create_magnetic_selector(kind, component) - else: - return self._create_nonmagnetic_selector(kind) - - def _create_magnetic_selector(self, kind, component): - _raise_error_if_kind_incorrect(kind, ("total", "xc")) - _raise_error_if_component_selected(component) - potential = self._get_potential(kind) - _raise_error_if_nonpolarized_potential(potential) - return index.Selector(maps={}, data=potential, reduction=_PotentialReduction) - - def _create_nonmagnetic_selector(self, kind): - potential = self._get_potential(kind) - maps = {0: self._create_map(potential)} - return index.Selector(maps, potential, reduction=_PotentialReduction) - - def _get_potential(self, kind): - return getattr(self._raw_potential, f"{kind}_potential") - - def _create_map(self, potential): - if _is_nonpolarized(potential): - return {choice: 0 for choice in _COMPONENTS[0]} - elif _is_collinear(potential): - return { - **{choice: 0 for choice in _COMPONENTS[0]}, - **{choice: 1 for choice in _COMPONENTS[3]}, - **{"up": slice(None), "down": slice(None)}, - } - return { - choice: component - for component, choices in _COMPONENTS.items() - for choice in choices - } - - def _create_isosurface(self, isolevel=0, color=None, opacity=0.6): - color = color or _config.VASP_COLORS["cyan"] - return view.Isosurface(isolevel, color, opacity) - - def _generate_items(self, kind): - potential = self._get_potential(kind) - if check.is_none(potential): - return - potential = np.moveaxis(potential, 0, -1).T - yield kind, potential[0] - if _is_collinear(potential): - yield f"{kind}_up", potential[0] + potential[1] - yield f"{kind}_down", potential[0] - potential[1] - elif _is_noncollinear(potential): - yield f"{kind}_magnetization", potential[1:] - - -@quantity("potential") -class Potential(view.Mixin): - """The local potential describes the interactions between electrons and ions. - - In DFT calculations, the local potential consists of various contributions, each - representing different aspects of the electron-electron and electron-ion - interactions. The ionic potential arises from the attraction between electrons and - the atomic nuclei. The Hartree potential accounts for the repulsion between - electrons resulting from the electron density itself. Additionally, the - exchange-correlation (xc) potential approximates the effects of electron exchange - and correlation. The accuracy of this approximation directly influences the - accuracy of the calculated properties. - - In VASP, the local potential is defined in real space on the FFT grid. You control - which potentials are written with the :tag:`WRT_POTENTIAL` tag. This class provides - the methods to read and visualize the potential. If you are interested in the - average potential, you may also look at the - :data:`~py4vasp.calculation.workfunction`. - - Examples - -------- - First, we create some example data do that you can follow along. Please define a - variable `path` with the path to a directory that exists and does not contain any - VASP calculation data. Alternatively, you can use your own data if you have run - VASP and construct `calculation` from it. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - See some basic information about the Potential by printing the object: - - >>> print(calculation.potential) - nonpolarized potential:... - - For your own postprocessing, you can read the potential data into a Python dictionary: - - >>> calculation.potential.read() - {'structure': {...}, 'total': array([[[...]]], ...), 'ionic': array([[[...]]], ...), 'xc': array([[[...]]], ...), 'hartree': array([[[...]]], ...)} - - You can also plot the 3d isosurface of the selected potential: - - >>> calculation.potential.plot() - View(...) - - Alternatively, you can visualize a contour plot of the potential in a plane: - - >>> calculation.potential.to_contour(c=0) - Graph(series=[Contour(data=array([[...]]...), ..., cut='c', ...)], ...) - - You can check possible selections for the potential: - - >>> calculation.potential.selections() - {'potential': ['default'...]...} - - Please check the documentation of each of these methods for more details on - how to use them and which options they provide. - """ - - def __init__(self, source, quantity_name="potential"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_potential): - return cls(source=DataSource(raw_potential)) - - def _handler_factory(self, raw): - return PotentialHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PotentialHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self) -> dict: - """Store all available contributions to the potential in a dictionary. - - Returns - ------- - dict - The dictionary contains the total potential as well as the potential - differences between up and down for collinear or the directional potential - for noncollinear calculations. If individual contributions to the potential - are available, these are returned, too. Structural information is given for - reference. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - PotentialHandler.to_dict, - ) - - def to_dict(self) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def to_view( - self, - selection: str = "total", - supercell: Optional[Union[int, np.ndarray]] = None, - **user_options, - ): - """Plot an isosurface of a selected potential. - - Parameters - ---------- - selection : str - Select the kind of potential of which you want the isosurface. - supercell : int or np.ndarray - If present the data is replicated the specified number of times along each - direction. - user_options - Further arguments with keyword that get directly passed on to the - visualizer. Most importantly, you can set isolevel (in eV) to adjust the - value at which the isosurface is drawn. - - Returns - ------- - View - A visualization of the potential isosurface within the crystal structure. - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PotentialHandler.to_view, - supercell=supercell, - **user_options, - ) - - @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) - def to_contour( - self, - selection: str = "total", - *, - a: Optional[float] = None, - b: Optional[float] = None, - c: Optional[float] = None, - normal: Optional[str] = None, - supercell: Optional[Union[int, np.ndarray]] = None, - ): - """Generate a 2D contour plot of the selected potential on a slice through the cell. - - {plane} - - Parameters - ---------- - selection : str, optional - Specifies which potential to plot. Can be any of the valid kinds - ("total", "ionic", "xc", "hartree"). For "total" and "xc" potentials - in spin-polarized calculations, you can select "up" or "down" components - (e.g., "total_up"). For noncollinear calculations, you can select - spin components ("x", "y", "z"). - - {parameters} - - Returns - ------- - Graph - A Graph object containing the contour plot. The plot shows the selected - potential component on the specified 2D slice. - - Examples - -------- - Cut a plane through the potential at the origin of the third lattice vector. - - >>> calculation.potential.to_contour(c=0) - - Plot the Hartree potential in the (100) plane crossing at 0.5 fractional coordinate - - >>> calculation.potential.to_contour(selection="hartree", a=0.5) - - Plot the sigma_z-component of the xc potential in a 2x2 supercell in the plane - defined by the first two lattice vectors. - - >>> calculation.potential.to_contour(selection="xc_z", c=0.2, supercell=2) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PotentialHandler.to_contour, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) - def to_quiver( - self, selection="total", *, a=None, b=None, c=None, normal=None, supercell=None - ): - """Generate a 2D quiver plot of the magnetic part of the potential on a slice. - - This method visualizes the vector field of the magnetization potential - (difference between spin-up and spin-down potentials for collinear cases, - or the vector components for noncollinear cases) as arrows on a 2D slice - through the simulation cell. - - {plane} - - Parameters - ---------- - selection : str, optional - Specifies which magnetic potential to plot. It must be a kind that - can have a magnetic component, i.e., "total" or "xc". - Component selection is not allowed for quiver plots as it inherently plots - the vector nature of the magnetization. - - {parameters} - - Returns - ------- - Graph - A Graph object containing the quiver plot. The plot shows - arrows representing the magnitude and direction of the magnetic - potential in the specified 2D slice. - - Examples - -------- - Plot the magnetization of the total potential in the a-b plane - - >>> calculation.potential.to_quiver(c=0) - - Plot the magnetization of the xc potential in the (010) plane - crossing at 0.25 fractional coordinate, using a 2x2 supercell. - - >>> calculation.potential.to_quiver(selection="xc", b=0.25, supercell=2) - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - PotentialHandler.to_quiver, - a=a, - b=b, - c=c, - normal=normal, - supercell=supercell, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - PotentialHandler.from_data, - PotentialHandler.to_database, - ) - - -class _PotentialReduction(index.Reduction): - def __init__(self, keys): - self._selection = keys[0] - - def __call__(self, array, axis): - if self._is_magnetic_potential(axis): - return np.moveaxis(array[1:], 0, -1) - if self._selection == "up": - return array[0] + array[1] - if self._selection == "down": - return array[0] - array[1] - return array[0] - - def _is_magnetic_potential(self, axis): - return axis == () - - -def _is_nonpolarized(potential): - return potential.shape[0] == 1 - - -def _is_collinear(potential): - return potential.shape[0] == 2 - - -def _is_noncollinear(potential): - return potential.shape[0] == 4 - - -def _raise_error_if_kind_incorrect(kind, valid_kinds=VALID_KINDS): - if kind in valid_kinds: - return - message = f"""\ -The selection "{kind}" is not a selection for the potential. Only the following \ -selections are allowed: "{'", "'.join(VALID_KINDS)}". \ -{suggest.did_you_mean(kind, valid_kinds)}Please check for spelling errors.""" - raise exception.IncorrectUsage(message) - - -def _raise_error_if_component_selected(component): - if not component: - return - message = f"Selecting a component {component} is not implemented for quiver plots." - raise exception.NotImplemented(message) - - -def _raise_error_if_no_data(data, kind="total"): - if data.is_none(): - message = f"Cannot find the {kind} potential data. " - if kind == "total": - message += ( - "Did you set LVTOT = T or WRT_POTENTIAL = total in the INCAR file?" - ) - else: - message += f"Did you set WRT_POTENTIAL = {kind} in the INCAR file?" - raise exception.NoData(message) - - -def _raise_error_if_nonpolarized_potential(potential): - if _is_nonpolarized(potential): - message = "Cannot visualize nonpolarized potential as quiver plot." - raise exception.DataMismatch(message) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy +import itertools +from typing import Optional, Union + +import numpy as np + +from py4vasp import _config, exception +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Potential_DB +from py4vasp._third_party import view +from py4vasp._util import ( + check, + density, + documentation, + index, + select, + slicing, + suggest, +) + +VALID_KINDS = ("total", "ionic", "xc", "hartree") +_COMPONENTS = { + 0: ["0", "unity", "sigma_0", "scalar"], + 1: ["1", "sigma_x", "x", "sigma_1"], + 2: ["2", "sigma_y", "y", "sigma_2"], + 3: ["3", "sigma_z", "z", "sigma_3"], +} +_COMMON_PARAMETERS = f"""\ +{slicing.PARAMETERS} +supercell : int or np.ndarray + Replicate the contour plot periodically a given number of times. If you + provide two different numbers, the resulting cell will be the two remaining + lattice vectors multiplied by the specific number. +""" + + +class PotentialHandler: + """Handler for potential data — performs all data access and transformation.""" + + def __init__(self, raw_potential: raw.Potential): + self._raw_potential = raw_potential + + @classmethod + def from_data(cls, raw_potential: raw.Potential) -> "PotentialHandler": + return cls(raw_potential) + + def __str__(self) -> str: + potential = self._raw_potential.total_potential + if _is_collinear(potential): + description = "collinear potential:" + elif _is_noncollinear(potential): + description = "noncollinear potential:" + else: + description = "nonpolarized potential:" + stoichiometry = _stoichiometry.Stoichiometry.from_data( + self._raw_potential.structure.stoichiometry + ) + structure = f"structure: {stoichiometry}" + grid = f"grid: {potential.shape[3]}, {potential.shape[2]}, {potential.shape[1]}" + available = "available: " + ", ".join( + kind for kind in VALID_KINDS if not self._get_potential(kind).is_none() + ) + return "\n ".join([description, structure, grid, available]) + + def to_dict(self) -> dict: + result = {"structure": self._structure().to_dict()} + items = [self._generate_items(kind) for kind in VALID_KINDS] + result.update(itertools.chain(*items)) + return result + + def to_database(self) -> Potential_DB: + total_potential_mean = None + total_potential_mean_up = None + total_potential_mean_down = None + total_potential_mean_magnetization = None + total_potential = self._get_potential("total") + if not check.is_none(total_potential): + total_potential = np.moveaxis(total_potential, 0, -1).T + total_potential_mean = float(np.mean(total_potential[0])) + total_potential_mean_up = ( + float(np.mean(total_potential[0] + total_potential[1])) + if _is_collinear(total_potential) + else total_potential_mean / 2.0 + ) + total_potential_mean_down = ( + float(np.mean(total_potential[0] - total_potential[1])) + if _is_collinear(total_potential) + else total_potential_mean / 2.0 + ) + total_potential_mean_magnetization = ( + float(np.mean(np.linalg.norm(total_potential[1:], axis=-1))) + if _is_noncollinear(total_potential) + else None + ) + + has_potential_dict = { + f"has_{kind}_potential": not check.is_none(self._get_potential(kind)) + for kind in VALID_KINDS + } + + return Potential_DB( + **has_potential_dict, + total_potential_mean=total_potential_mean, + total_potential_mean_up=total_potential_mean_up, + total_potential_mean_down=total_potential_mean_down, + total_potential_mean_magnetization=total_potential_mean_magnetization, + ) + + def to_view( + self, + selection: str = "total", + supercell: Optional[Union[int, np.ndarray]] = None, + **user_options, + ): + potentials = dict(self._get_potentials(selection)) + isosurface = self._create_isosurface(**user_options) + viewer = self._structure().to_view(supercell) + viewer.grid_scalars = [ + view.GridQuantity( + quantity=data[np.newaxis], + label=label, + isosurfaces=[isosurface], + ) + for label, data in potentials.items() + ] + return viewer + + def to_contour( + self, + selection: str = "total", + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ): + potentials = dict(self._get_potentials(selection)) + visualizer = density.Visualizer(self._structure()) + slice_arguments = density.SliceArguments(a, b, c, normal, supercell) + return visualizer.to_contour(potentials, slice_arguments, isolevels=False) + + def to_quiver( + self, selection="total", *, a=None, b=None, c=None, normal=None, supercell=None + ): + potentials = dict(self._get_potentials(selection, is_magnetic=True)) + visualizer = density.Visualizer(self._structure()) + slice_arguments = density.SliceArguments(a, b, c, normal, supercell) + return visualizer.to_quiver(potentials, slice_arguments) + + def _structure(self): + return StructureHandler.from_data(self._raw_potential.structure) + + def _get_potentials(self, selection, is_magnetic=False): + tree = select.Tree.from_selection(selection) + for selection in tree.selections(): + kind, component = self._determine_kind_and_component(selection) + selector = self._create_selector(kind, component, is_magnetic) + component_label = component[0] if component else "" + yield self._get_label(kind, component_label), selector[component].T + + def _determine_kind_and_component(self, selection): + for kind in VALID_KINDS: + if kind in selection: + remaining = list(selection) + remaining.remove(kind) + return kind, tuple(remaining) + return "total", selection + + def _get_label(self, kind, component): + return f"{kind} potential" + (f"({component})" if component else "") + + def _create_selector(self, kind, component, is_magnetic): + if is_magnetic: + return self._create_magnetic_selector(kind, component) + else: + return self._create_nonmagnetic_selector(kind) + + def _create_magnetic_selector(self, kind, component): + _raise_error_if_kind_incorrect(kind, ("total", "xc")) + _raise_error_if_component_selected(component) + potential = self._get_potential(kind) + _raise_error_if_nonpolarized_potential(potential) + return index.Selector(maps={}, data=potential, reduction=_PotentialReduction) + + def _create_nonmagnetic_selector(self, kind): + potential = self._get_potential(kind) + maps = {0: self._create_map(potential)} + return index.Selector(maps, potential, reduction=_PotentialReduction) + + def _get_potential(self, kind): + return getattr(self._raw_potential, f"{kind}_potential") + + def _create_map(self, potential): + if _is_nonpolarized(potential): + return {choice: 0 for choice in _COMPONENTS[0]} + elif _is_collinear(potential): + return { + **{choice: 0 for choice in _COMPONENTS[0]}, + **{choice: 1 for choice in _COMPONENTS[3]}, + **{"up": slice(None), "down": slice(None)}, + } + return { + choice: component + for component, choices in _COMPONENTS.items() + for choice in choices + } + + def _create_isosurface(self, isolevel=0, color=None, opacity=0.6): + color = color or _config.VASP_COLORS["cyan"] + return view.Isosurface(isolevel, color, opacity) + + def _generate_items(self, kind): + potential = self._get_potential(kind) + if check.is_none(potential): + return + potential = np.moveaxis(potential, 0, -1).T + yield kind, potential[0] + if _is_collinear(potential): + yield f"{kind}_up", potential[0] + potential[1] + yield f"{kind}_down", potential[0] - potential[1] + elif _is_noncollinear(potential): + yield f"{kind}_magnetization", potential[1:] + + +@quantity("potential") +class Potential(view.Mixin): + """The local potential describes the interactions between electrons and ions. + + In DFT calculations, the local potential consists of various contributions, each + representing different aspects of the electron-electron and electron-ion + interactions. The ionic potential arises from the attraction between electrons and + the atomic nuclei. The Hartree potential accounts for the repulsion between + electrons resulting from the electron density itself. Additionally, the + exchange-correlation (xc) potential approximates the effects of electron exchange + and correlation. The accuracy of this approximation directly influences the + accuracy of the calculated properties. + + In VASP, the local potential is defined in real space on the FFT grid. You control + which potentials are written with the :tag:`WRT_POTENTIAL` tag. This class provides + the methods to read and visualize the potential. If you are interested in the + average potential, you may also look at the + :data:`~py4vasp.calculation.workfunction`. + + Examples + -------- + First, we create some example data do that you can follow along. Please define a + variable `path` with the path to a directory that exists and does not contain any + VASP calculation data. Alternatively, you can use your own data if you have run + VASP and construct `calculation` from it. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + See some basic information about the Potential by printing the object: + + >>> print(calculation.potential) + nonpolarized potential:... + + For your own postprocessing, you can read the potential data into a Python dictionary: + + >>> calculation.potential.read() + {'structure': {...}, 'total': array([[[...]]], ...), 'ionic': array([[[...]]], ...), 'xc': array([[[...]]], ...), 'hartree': array([[[...]]], ...)} + + You can also plot the 3d isosurface of the selected potential: + + >>> calculation.potential.plot() + View(...) + + Alternatively, you can visualize a contour plot of the potential in a plane: + + >>> calculation.potential.to_contour(c=0) + Graph(series=[Contour(data=array([[...]]...), ..., cut='c', ...)], ...) + + You can check possible selections for the potential: + + >>> calculation.potential.selections() + {'potential': ['default'...]...} + + Please check the documentation of each of these methods for more details on + how to use them and which options they provide. + """ + + def __init__(self, source, quantity_name="potential"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_potential): + return cls(source=DataSource(raw_potential)) + + def _handler_factory(self, raw): + return PotentialHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PotentialHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self) -> dict: + """Store all available contributions to the potential in a dictionary. + + Returns + ------- + dict + The dictionary contains the total potential as well as the potential + differences between up and down for collinear or the directional potential + for noncollinear calculations. If individual contributions to the potential + are available, these are returned, too. Structural information is given for + reference. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + PotentialHandler.to_dict, + ) + + def to_dict(self) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def to_view( + self, + selection: str = "total", + supercell: Optional[Union[int, np.ndarray]] = None, + **user_options, + ): + """Plot an isosurface of a selected potential. + + Parameters + ---------- + selection : str + Select the kind of potential of which you want the isosurface. + supercell : int or np.ndarray + If present the data is replicated the specified number of times along each + direction. + user_options + Further arguments with keyword that get directly passed on to the + visualizer. Most importantly, you can set isolevel (in eV) to adjust the + value at which the isosurface is drawn. + + Returns + ------- + View + A visualization of the potential isosurface within the crystal structure. + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PotentialHandler.to_view, + supercell=supercell, + **user_options, + ) + + @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) + def to_contour( + self, + selection: str = "total", + *, + a: Optional[float] = None, + b: Optional[float] = None, + c: Optional[float] = None, + normal: Optional[str] = None, + supercell: Optional[Union[int, np.ndarray]] = None, + ): + """Generate a 2D contour plot of the selected potential on a slice through the cell. + + {plane} + + Parameters + ---------- + selection : str, optional + Specifies which potential to plot. Can be any of the valid kinds + ("total", "ionic", "xc", "hartree"). For "total" and "xc" potentials + in spin-polarized calculations, you can select "up" or "down" components + (e.g., "total_up"). For noncollinear calculations, you can select + spin components ("x", "y", "z"). + + {parameters} + + Returns + ------- + Graph + A Graph object containing the contour plot. The plot shows the selected + potential component on the specified 2D slice. + + Examples + -------- + Cut a plane through the potential at the origin of the third lattice vector. + + >>> calculation.potential.to_contour(c=0) + + Plot the Hartree potential in the (100) plane crossing at 0.5 fractional coordinate + + >>> calculation.potential.to_contour(selection="hartree", a=0.5) + + Plot the sigma_z-component of the xc potential in a 2x2 supercell in the plane + defined by the first two lattice vectors. + + >>> calculation.potential.to_contour(selection="xc_z", c=0.2, supercell=2) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PotentialHandler.to_contour, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + @documentation.format(plane=slicing.PLANE, parameters=_COMMON_PARAMETERS) + def to_quiver( + self, selection="total", *, a=None, b=None, c=None, normal=None, supercell=None + ): + """Generate a 2D quiver plot of the magnetic part of the potential on a slice. + + This method visualizes the vector field of the magnetization potential + (difference between spin-up and spin-down potentials for collinear cases, + or the vector components for noncollinear cases) as arrows on a 2D slice + through the simulation cell. + + {plane} + + Parameters + ---------- + selection : str, optional + Specifies which magnetic potential to plot. It must be a kind that + can have a magnetic component, i.e., "total" or "xc". + Component selection is not allowed for quiver plots as it inherently plots + the vector nature of the magnetization. + + {parameters} + + Returns + ------- + Graph + A Graph object containing the quiver plot. The plot shows + arrows representing the magnitude and direction of the magnetic + potential in the specified 2D slice. + + Examples + -------- + Plot the magnetization of the total potential in the a-b plane + + >>> calculation.potential.to_quiver(c=0) + + Plot the magnetization of the xc potential in the (010) plane + crossing at 0.25 fractional coordinate, using a 2x2 supercell. + + >>> calculation.potential.to_quiver(selection="xc", b=0.25, supercell=2) + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + PotentialHandler.to_quiver, + a=a, + b=b, + c=c, + normal=normal, + supercell=supercell, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + PotentialHandler.from_data, + PotentialHandler.to_database, + ) + + +class _PotentialReduction(index.Reduction): + def __init__(self, keys): + self._selection = keys[0] + + def __call__(self, array, axis): + if self._is_magnetic_potential(axis): + return np.moveaxis(array[1:], 0, -1) + if self._selection == "up": + return array[0] + array[1] + if self._selection == "down": + return array[0] - array[1] + return array[0] + + def _is_magnetic_potential(self, axis): + return axis == () + + +def _is_nonpolarized(potential): + return potential.shape[0] == 1 + + +def _is_collinear(potential): + return potential.shape[0] == 2 + + +def _is_noncollinear(potential): + return potential.shape[0] == 4 + + +def _raise_error_if_kind_incorrect(kind, valid_kinds=VALID_KINDS): + if kind in valid_kinds: + return + message = f"""\ +The selection "{kind}" is not a selection for the potential. Only the following \ +selections are allowed: "{'", "'.join(VALID_KINDS)}". \ +{suggest.did_you_mean(kind, valid_kinds)}Please check for spelling errors.""" + raise exception.IncorrectUsage(message) + + +def _raise_error_if_component_selected(component): + if not component: + return + message = f"Selecting a component {component} is not implemented for quiver plots." + raise exception.NotImplemented(message) + + +def _raise_error_if_no_data(data, kind="total"): + if data.is_none(): + message = f"Cannot find the {kind} potential data. " + if kind == "total": + message += ( + "Did you set LVTOT = T or WRT_POTENTIAL = total in the INCAR file?" + ) + else: + message += f"Did you set WRT_POTENTIAL = {kind} in the INCAR file?" + raise exception.NoData(message) + + +def _raise_error_if_nonpolarized_potential(potential): + if _is_nonpolarized(potential): + message = "Cannot visualize nonpolarized potential as quiver plot." + raise exception.DataMismatch(message) diff --git a/src/py4vasp/_calculation/projector.py b/src/py4vasp/_calculation/projector.py index 997eee92d..13f22434a 100644 --- a/src/py4vasp/_calculation/projector.py +++ b/src/py4vasp/_calculation/projector.py @@ -1,387 +1,387 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from py4vasp import exception -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw import data as raw -from py4vasp._raw.data_db import Projector_DB -from py4vasp._util import check, convert, index, select - -SPIN_PROJECTION = "is_spin_projection" -selection_doc = """\ -selection : str - A string specifying the projection of the orbitals. There are four distinct - possibilities: - - - To specify the **atom**, you can either use its element name (Si, Al, ...) - or its index as given in the input file (1, 2, ...). For the latter - option it is also possible to specify ranges (e.g. 1:4). - - To select a particular **orbital** you can give a string (s, px, dxz, ...) - or select multiple orbitals by their angular momentum (s, p, d, f). - - For the **spin**, you have the options up, down, or total. - - If you used a different **k**-point mesh choose "kpoints_opt" or "kpoints_wan" - to select them instead of the default mesh specified in the KPOINTS file. - - You separate multiple selections by commas or whitespace and can nest them using - parenthesis, e.g. `Sr(s, p)` or `s(up), p(down)`. The order of the selections - does not matter, but it is case sensitive to distinguish p (angular momentum - l = 1) from P (phosphorus). - - It is possible to add or subtract different components, e.g., a selection of - "Ti(d) - O(p)" would project onto the d orbitals of Ti and the p orbitals of O and - then compute the difference of these two selections. - - If you are unsure about the specific projections that are available, you can use - - >>> calculation.projector.selections() - {\'atom\': [...], \'orbital\': [...], \'spin\': [...]} - - to get a list of all available ones. -""" - - -class ProjectorHandler: - """Handler for projector data.""" - - _missing_data_message = "No projectors found, please verify the LORBIT tag is set." - - def __init__(self, raw_projector: raw.Projector): - self._raw_projector = raw_projector - - @classmethod - def from_data(cls, raw_projector: raw.Projector) -> "ProjectorHandler": - return cls(raw_projector) - - def __str__(self): - if self._raw_projector.orbital_types.is_none(): - return "no projectors" - if self._is_collinear: - spin_projection = "\n spin: total, up, down" - elif self._is_noncollinear: - spin_projection = "\n spin: total, sigma_x, sigma_y, sigma_z" - else: - spin_projection = "" - return f"""projectors: - atoms: {", ".join(self._stoichiometry().ion_types())} - orbitals: {", ".join(self._orbital_types())}""" + spin_projection - - def to_dict(self) -> dict: - """Return a map from labels to indices in the arrays produced by VASP. - - Returns - ------- - dict - A dictionary containing three dictionaries for spin, atom, and orbitals. - Each of those describes which indices VASP uses to store certain elements - for projected quantities. If VASP was run without setting :tag:`LORBIT` - this will return an empty dictionary. - - Examples - -------- - - For nonpolarized Fe3O4 with :tag:`LORBIT` = 10, this would work like this - - >>> import pprint - >>> from py4vasp import demo - >>> calculation = demo.calculation(path, selection="collinear") - >>> pprint.pp(calculation.projector.to_dict()) - {'atom': {'Fe': slice(0, 3, None), - '1': slice(0, 1, None), - '2': slice(1, 2, None), - '3': slice(2, 3, None), - 'O': slice(3, 7, None), - '4': slice(3, 4, None), - '5': slice(4, 5, None), - '6': slice(5, 6, None), - '7': slice(6, 7, None)}, - 'orbital': {'s': slice(0, 1, None), - 'p': slice(1, 2, None), - 'd': slice(2, 3, None), - 'f': slice(3, 4, None)}, - 'spin': {'total': slice(0, 2, None), - 'up': slice(0, 1, None), - 'down': slice(1, 2, None)}} - """ - if self._raw_projector.orbital_types.is_none(): - return {} - atom_dict = self._init_atom_dict() - orbital_dict = self._init_orbital_dict() - spin_dict = self._init_spin_dict() - return {"atom": atom_dict, "orbital": orbital_dict, "spin": spin_dict} - - def to_database(self) -> dict: - return Projector_DB( - orbital_types=( - sorted(list(self._init_orbital_dict().keys()), key=self._sort_key) - if not check.is_none(self._raw_projector.orbital_types) - else None - ) - ) - - def selections(self) -> dict: - dicts = self.to_dict() - if len(dicts) == 0: - return dicts - return { - "atom": sorted(dicts["atom"], key=self._sort_key), - "orbital": sorted(dicts["orbital"], key=self._sort_key), - "spin": sorted(dicts["spin"], key=self._sort_key), - } - - def project(self, selection, projections): - if not selection: - return {} - self._raise_error_if_orbitals_missing() - selector = self._make_selector(projections) - return dict(self._create_projections(selector, selection)) - - def _stoichiometry(self): - return _stoichiometry.Stoichiometry.from_data(self._raw_projector.stoichiometry) - - def _orbital_types(self): - clean_string = lambda orbital: convert.text_to_string(orbital).strip() - for orbital in self._raw_projector.orbital_types: - orbital = clean_string(orbital) - if orbital == "x2-y2": - yield "dx2y2" - else: - yield orbital - - def _init_atom_dict(self): - return { - key: value.indices - for key, value in self._stoichiometry().read().items() - if key != select.all - } - - def _init_orbital_dict(self): - orbital_dict = { - orbital: slice(i, i + 1) for i, orbital in enumerate(self._orbital_types()) - } - if "px" in orbital_dict: - orbital_dict["p"] = slice(1, 4) - orbital_dict["d"] = slice(4, 9) - orbital_dict["f"] = slice(9, 16) - return orbital_dict - - def _init_spin_dict(self): - if self._is_nonpolarized: - return {"total": slice(0, 1)} - if self._is_collinear: - return {"total": slice(0, 2), "up": slice(0, 1), "down": slice(1, 2)} - return { - "total": slice(0, 1), - "sigma_x": slice(1, 2), - "sigma_y": slice(2, 3), - "sigma_z": slice(3, 4), - "x": slice(1, 2), - "y": slice(2, 3), - "z": slice(3, 4), - "sigma_1": slice(1, 2), - "sigma_2": slice(2, 3), - "sigma_3": slice(3, 4), - } - - @property - def _is_nonpolarized(self): - return self._raw_projector.number_spin_projections == 1 - - @property - def _is_collinear(self): - return self._raw_projector.number_spin_projections == 2 - - @property - def _is_noncollinear(self): - return self._raw_projector.number_spin_projections == 4 - - def _sort_key(self, key): - spin_keys = [ - "total", - "up", - "down", - "sigma_x", - "sigma_y", - "sigma_z", - "x", - "y", - "z", - "sigma_1", - "sigma_2", - "sigma_3", - ] - orbital_keys = ["s", "p", "d", "f"] - if key in spin_keys: - return 0 - if key[:1] in orbital_keys: - return str(orbital_keys.index(key[:1])) + key - if key.isdecimal(): - return int(key) - assert key.istitle() - return 0 - - def _make_selector(self, projections): - maps = self.to_dict() - maps = {1: maps["atom"], 2: maps["orbital"], 0: maps["spin"]} - try: - return index.Selector(maps, projections, use_number_labels=True) - except exception._Py4VaspInternalError: - message = f"""Error reading the projections. Please make sure that the passed - projections has the right format, i.e., the indices correspond to spin, - atom, and orbital, respectively.""" - raise exception.IncorrectUsage(message) from None - - def _create_projections(self, selector, selection): - spin_projections = [] - tree = select.Tree.from_selection(selection) - for selection in tree.selections(): - if self._is_nonpolarized or self._spin_selected(selection): - label, weight = self._create_projection(selector, selection) - if "total" not in selection: - spin_projections.append(label) - yield label, weight - elif self._is_collinear: - yield self._create_projection(selector, selection + ("up",)) - yield self._create_projection(selector, selection + ("down",)) - else: - yield self._create_projection(selector, selection + ("total",)) - if self._is_noncollinear: - yield SPIN_PROJECTION, spin_projections - - def _create_projection(self, selector, selection): - label = selector.label(selection) - if self._is_noncollinear: - label = label.removesuffix("_total") - return label, selector[selection] - - def _spin_selected(self, selection): - return any( - select.contains(selection, choice) for choice in self._init_spin_dict() - ) - - def _raise_error_if_orbitals_missing(self): - if self._raw_projector.orbital_types.is_none(): - message = "Projectors are not available, rerun VASP setting LORBIT >= 10." - raise exception.IncorrectUsage(message) - - -@quantity("projector") -class Projector: - """The projectors used for atom and orbital resolved quantities. - - This is a utility class that facilitates projecting quantities such as the - electronic band structure and the DOS on atoms and orbitals. As a user, you can - investigate the available projections with the :meth:`read` or :meth:`selections` - methods. The former is useful for scripts, when you need to know which array - index corresponds to which orbital or atom. The latter describes the available - selections that you can use in the methods that project on orbitals or atoms. - """ - - def __init__(self, source, quantity_name="projector"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_projector): - return cls(source=DataSource(raw_projector)) - - def _handler_factory(self, raw): - return ProjectorHandler.from_data(raw) - - def __str__(self, selection=None): - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ProjectorHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def read(self, selection=None) -> dict: - """Return a map from labels to indices in the arrays produced by VASP. - - Returns - ------- - dict - A dictionary containing three dictionaries for spin, atom, and orbitals. - Each of those describes which indices VASP uses to store certain elements - for projected quantities. If VASP was run without setting :tag:`LORBIT` - this will return an empty dictionary. - - Examples - -------- - - For nonpolarized Fe3O4 with :tag:`LORBIT` = 10, this would work like this - - >>> import pprint - >>> from py4vasp import demo - >>> calculation = demo.calculation(path, selection="collinear") - >>> pprint.pp(calculation.projector.read()) - {'atom': {'Fe': slice(0, 3, None), - '1': slice(0, 1, None), - '2': slice(1, 2, None), - '3': slice(2, 3, None), - 'O': slice(3, 7, None), - '4': slice(3, 4, None), - '5': slice(4, 5, None), - '6': slice(5, 6, None), - '7': slice(6, 7, None)}, - 'orbital': {'s': slice(0, 1, None), - 'p': slice(1, 2, None), - 'd': slice(2, 3, None), - 'f': slice(3, 4, None)}, - 'spin': {'total': slice(0, 2, None), - 'up': slice(0, 1, None), - 'down': slice(1, 2, None)}} - """ - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ProjectorHandler.to_dict, - ) - - def to_dict(self, selection=None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(selection=selection) - - def selections(self, selection=None) -> dict: - from py4vasp._raw import definition as raw_module - - handler_selections = merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ProjectorHandler.selections, - ) - return handler_selections - - def project(self, selection, projections): - return merge_default( - self._source, - self._quantity_name, - selection, - self._handler_factory, - ProjectorHandler.project, - projections, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - ProjectorHandler.from_data, - ProjectorHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from py4vasp import exception +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw import data as raw +from py4vasp._raw.data_db import Projector_DB +from py4vasp._util import check, convert, index, select + +SPIN_PROJECTION = "is_spin_projection" +selection_doc = """\ +selection : str + A string specifying the projection of the orbitals. There are four distinct + possibilities: + + - To specify the **atom**, you can either use its element name (Si, Al, ...) + or its index as given in the input file (1, 2, ...). For the latter + option it is also possible to specify ranges (e.g. 1:4). + - To select a particular **orbital** you can give a string (s, px, dxz, ...) + or select multiple orbitals by their angular momentum (s, p, d, f). + - For the **spin**, you have the options up, down, or total. + - If you used a different **k**-point mesh choose "kpoints_opt" or "kpoints_wan" + to select them instead of the default mesh specified in the KPOINTS file. + + You separate multiple selections by commas or whitespace and can nest them using + parenthesis, e.g. `Sr(s, p)` or `s(up), p(down)`. The order of the selections + does not matter, but it is case sensitive to distinguish p (angular momentum + l = 1) from P (phosphorus). + + It is possible to add or subtract different components, e.g., a selection of + "Ti(d) - O(p)" would project onto the d orbitals of Ti and the p orbitals of O and + then compute the difference of these two selections. + + If you are unsure about the specific projections that are available, you can use + + >>> calculation.projector.selections() + {\'atom\': [...], \'orbital\': [...], \'spin\': [...]} + + to get a list of all available ones. +""" + + +class ProjectorHandler: + """Handler for projector data.""" + + _missing_data_message = "No projectors found, please verify the LORBIT tag is set." + + def __init__(self, raw_projector: raw.Projector): + self._raw_projector = raw_projector + + @classmethod + def from_data(cls, raw_projector: raw.Projector) -> "ProjectorHandler": + return cls(raw_projector) + + def __str__(self): + if self._raw_projector.orbital_types.is_none(): + return "no projectors" + if self._is_collinear: + spin_projection = "\n spin: total, up, down" + elif self._is_noncollinear: + spin_projection = "\n spin: total, sigma_x, sigma_y, sigma_z" + else: + spin_projection = "" + return f"""projectors: + atoms: {", ".join(self._stoichiometry().ion_types())} + orbitals: {", ".join(self._orbital_types())}""" + spin_projection + + def to_dict(self) -> dict: + """Return a map from labels to indices in the arrays produced by VASP. + + Returns + ------- + dict + A dictionary containing three dictionaries for spin, atom, and orbitals. + Each of those describes which indices VASP uses to store certain elements + for projected quantities. If VASP was run without setting :tag:`LORBIT` + this will return an empty dictionary. + + Examples + -------- + + For nonpolarized Fe3O4 with :tag:`LORBIT` = 10, this would work like this + + >>> import pprint + >>> from py4vasp import demo + >>> calculation = demo.calculation(path, selection="collinear") + >>> pprint.pp(calculation.projector.to_dict()) + {'atom': {'Fe': slice(0, 3, None), + '1': slice(0, 1, None), + '2': slice(1, 2, None), + '3': slice(2, 3, None), + 'O': slice(3, 7, None), + '4': slice(3, 4, None), + '5': slice(4, 5, None), + '6': slice(5, 6, None), + '7': slice(6, 7, None)}, + 'orbital': {'s': slice(0, 1, None), + 'p': slice(1, 2, None), + 'd': slice(2, 3, None), + 'f': slice(3, 4, None)}, + 'spin': {'total': slice(0, 2, None), + 'up': slice(0, 1, None), + 'down': slice(1, 2, None)}} + """ + if self._raw_projector.orbital_types.is_none(): + return {} + atom_dict = self._init_atom_dict() + orbital_dict = self._init_orbital_dict() + spin_dict = self._init_spin_dict() + return {"atom": atom_dict, "orbital": orbital_dict, "spin": spin_dict} + + def to_database(self) -> dict: + return Projector_DB( + orbital_types=( + sorted(list(self._init_orbital_dict().keys()), key=self._sort_key) + if not check.is_none(self._raw_projector.orbital_types) + else None + ) + ) + + def selections(self) -> dict: + dicts = self.to_dict() + if len(dicts) == 0: + return dicts + return { + "atom": sorted(dicts["atom"], key=self._sort_key), + "orbital": sorted(dicts["orbital"], key=self._sort_key), + "spin": sorted(dicts["spin"], key=self._sort_key), + } + + def project(self, selection, projections): + if not selection: + return {} + self._raise_error_if_orbitals_missing() + selector = self._make_selector(projections) + return dict(self._create_projections(selector, selection)) + + def _stoichiometry(self): + return _stoichiometry.Stoichiometry.from_data(self._raw_projector.stoichiometry) + + def _orbital_types(self): + clean_string = lambda orbital: convert.text_to_string(orbital).strip() + for orbital in self._raw_projector.orbital_types: + orbital = clean_string(orbital) + if orbital == "x2-y2": + yield "dx2y2" + else: + yield orbital + + def _init_atom_dict(self): + return { + key: value.indices + for key, value in self._stoichiometry().read().items() + if key != select.all + } + + def _init_orbital_dict(self): + orbital_dict = { + orbital: slice(i, i + 1) for i, orbital in enumerate(self._orbital_types()) + } + if "px" in orbital_dict: + orbital_dict["p"] = slice(1, 4) + orbital_dict["d"] = slice(4, 9) + orbital_dict["f"] = slice(9, 16) + return orbital_dict + + def _init_spin_dict(self): + if self._is_nonpolarized: + return {"total": slice(0, 1)} + if self._is_collinear: + return {"total": slice(0, 2), "up": slice(0, 1), "down": slice(1, 2)} + return { + "total": slice(0, 1), + "sigma_x": slice(1, 2), + "sigma_y": slice(2, 3), + "sigma_z": slice(3, 4), + "x": slice(1, 2), + "y": slice(2, 3), + "z": slice(3, 4), + "sigma_1": slice(1, 2), + "sigma_2": slice(2, 3), + "sigma_3": slice(3, 4), + } + + @property + def _is_nonpolarized(self): + return self._raw_projector.number_spin_projections == 1 + + @property + def _is_collinear(self): + return self._raw_projector.number_spin_projections == 2 + + @property + def _is_noncollinear(self): + return self._raw_projector.number_spin_projections == 4 + + def _sort_key(self, key): + spin_keys = [ + "total", + "up", + "down", + "sigma_x", + "sigma_y", + "sigma_z", + "x", + "y", + "z", + "sigma_1", + "sigma_2", + "sigma_3", + ] + orbital_keys = ["s", "p", "d", "f"] + if key in spin_keys: + return 0 + if key[:1] in orbital_keys: + return str(orbital_keys.index(key[:1])) + key + if key.isdecimal(): + return int(key) + assert key.istitle() + return 0 + + def _make_selector(self, projections): + maps = self.to_dict() + maps = {1: maps["atom"], 2: maps["orbital"], 0: maps["spin"]} + try: + return index.Selector(maps, projections, use_number_labels=True) + except exception._Py4VaspInternalError: + message = f"""Error reading the projections. Please make sure that the passed + projections has the right format, i.e., the indices correspond to spin, + atom, and orbital, respectively.""" + raise exception.IncorrectUsage(message) from None + + def _create_projections(self, selector, selection): + spin_projections = [] + tree = select.Tree.from_selection(selection) + for selection in tree.selections(): + if self._is_nonpolarized or self._spin_selected(selection): + label, weight = self._create_projection(selector, selection) + if "total" not in selection: + spin_projections.append(label) + yield label, weight + elif self._is_collinear: + yield self._create_projection(selector, selection + ("up",)) + yield self._create_projection(selector, selection + ("down",)) + else: + yield self._create_projection(selector, selection + ("total",)) + if self._is_noncollinear: + yield SPIN_PROJECTION, spin_projections + + def _create_projection(self, selector, selection): + label = selector.label(selection) + if self._is_noncollinear: + label = label.removesuffix("_total") + return label, selector[selection] + + def _spin_selected(self, selection): + return any( + select.contains(selection, choice) for choice in self._init_spin_dict() + ) + + def _raise_error_if_orbitals_missing(self): + if self._raw_projector.orbital_types.is_none(): + message = "Projectors are not available, rerun VASP setting LORBIT >= 10." + raise exception.IncorrectUsage(message) + + +@quantity("projector") +class Projector: + """The projectors used for atom and orbital resolved quantities. + + This is a utility class that facilitates projecting quantities such as the + electronic band structure and the DOS on atoms and orbitals. As a user, you can + investigate the available projections with the :meth:`read` or :meth:`selections` + methods. The former is useful for scripts, when you need to know which array + index corresponds to which orbital or atom. The latter describes the available + selections that you can use in the methods that project on orbitals or atoms. + """ + + def __init__(self, source, quantity_name="projector"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_projector): + return cls(source=DataSource(raw_projector)) + + def _handler_factory(self, raw): + return ProjectorHandler.from_data(raw) + + def __str__(self, selection=None): + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ProjectorHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def read(self, selection=None) -> dict: + """Return a map from labels to indices in the arrays produced by VASP. + + Returns + ------- + dict + A dictionary containing three dictionaries for spin, atom, and orbitals. + Each of those describes which indices VASP uses to store certain elements + for projected quantities. If VASP was run without setting :tag:`LORBIT` + this will return an empty dictionary. + + Examples + -------- + + For nonpolarized Fe3O4 with :tag:`LORBIT` = 10, this would work like this + + >>> import pprint + >>> from py4vasp import demo + >>> calculation = demo.calculation(path, selection="collinear") + >>> pprint.pp(calculation.projector.read()) + {'atom': {'Fe': slice(0, 3, None), + '1': slice(0, 1, None), + '2': slice(1, 2, None), + '3': slice(2, 3, None), + 'O': slice(3, 7, None), + '4': slice(3, 4, None), + '5': slice(4, 5, None), + '6': slice(5, 6, None), + '7': slice(6, 7, None)}, + 'orbital': {'s': slice(0, 1, None), + 'p': slice(1, 2, None), + 'd': slice(2, 3, None), + 'f': slice(3, 4, None)}, + 'spin': {'total': slice(0, 2, None), + 'up': slice(0, 1, None), + 'down': slice(1, 2, None)}} + """ + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ProjectorHandler.to_dict, + ) + + def to_dict(self, selection=None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(selection=selection) + + def selections(self, selection=None) -> dict: + from py4vasp._raw import definition as raw_module + + handler_selections = merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ProjectorHandler.selections, + ) + return handler_selections + + def project(self, selection, projections): + return merge_default( + self._source, + self._quantity_name, + selection, + self._handler_factory, + ProjectorHandler.project, + projections, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + ProjectorHandler.from_data, + ProjectorHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/run_info.py b/src/py4vasp/_calculation/run_info.py index 2497b998d..4ef746115 100644 --- a/src/py4vasp/_calculation/run_info.py +++ b/src/py4vasp/_calculation/run_info.py @@ -1,220 +1,220 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from contextlib import suppress - -from py4vasp import raw -from py4vasp._calculation import bandgap as bandgap_module -from py4vasp._calculation import exception -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import RunInfo_DB -from py4vasp._raw.data_wrapper import VaspData -from py4vasp._util import check - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - exception.OutdatedVaspVersion, - exception.NoData, - AttributeError, - TypeError, - ValueError, -) - - -class RunInfoHandler: - """Handler for run info. Works with exactly one raw.RunInfo object.""" - - def __init__(self, raw_run_info: raw.RunInfo): - self._raw_run_info = raw_run_info - - @classmethod - def from_data(cls, raw_run_info: raw.RunInfo) -> "RunInfoHandler": - return cls(raw_run_info) - - def to_dict(self) -> dict: - """Convert the run information to a dictionary.""" - return { - **self._dict_from_runtime(), - **self._dict_from_system(), - **self._dict_from_structure(), - **self._dict_additional_collection(), - **self._dict_from_contcar(), - **self._dict_from_phonon_dispersion(), - } - - def to_database(self) -> dict: - """Serialize run info for the database.""" - return RunInfo_DB(**self.to_dict()) - - def _read_attr(self, *keys: str): - data = self._raw_run_info - for key in keys: - data = getattr(data, key, None) - if data is None: - return None - try: - is_none = data.is_none() or data is None - except AttributeError: - is_none = False - return data[:] if not is_none else None - - def _dict_additional_collection(self) -> dict: - fermi_energy = None - with suppress(exception.NoData): - fermi_energy = self._raw_run_info.fermi_energy - if isinstance(fermi_energy, VaspData): - fermi_energy = fermi_energy._data - - is_success = None # TODO implement - - is_collinear = self._is_collinear() - is_noncollinear = self._is_noncollinear() - is_metallic = self._is_metallic() - is_magnetic = None # TODO implement - magnetic_order = None # TODO implement - - grid_coarse_shape = ( - None # TODO implement for FFT grid (currently not written to H5) - ) - grid_fine_shape = ( - None # TODO implement for FFT grid (currently not written to H5) - ) - - return { - "grid_coarse_shape": grid_coarse_shape, - "grid_fine_shape": grid_fine_shape, - "is_success": is_success, - "fermi_energy": fermi_energy, - "is_collinear": is_collinear, - "is_noncollinear": is_noncollinear, - "is_metallic": is_metallic, - "is_magnetic": is_magnetic, - "magnetization_order": magnetic_order, - } - - def _is_collinear(self): - if not check.is_none(self._raw_run_info.len_dos): - return self._raw_run_info.len_dos == 2 - else: - if not check.is_none(self._raw_run_info.band_dispersion_eigenvalues): - return len(self._raw_run_info.band_dispersion_eigenvalues) == 2 - else: - return None - - def _is_noncollinear(self): - if not check.is_none(self._raw_run_info.len_dos): - return self._raw_run_info.len_dos == 4 - else: - if not check.is_none(self._raw_run_info.band_projections): - return len(self._raw_run_info.band_projections) == 4 - else: - return None - - def _is_metallic(self): - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - if check.is_none(self._raw_run_info.bandgap): - return None - gap = bandgap_module.BandgapHandler.from_data(self._raw_run_info.bandgap) - return all(gap._output_gap("fundamental", to_string=False) <= 0.0) - return None - - def _dict_from_system(self) -> dict: - system_tag = None - with suppress(exception.NoData): - system_tag = self._read_attr("system", "system") - return { - "system_tag": system_tag, - } - - def _dict_from_runtime(self) -> dict: - vasp_version = None - with suppress(exception.NoData): - runtime_data = self._raw_run_info.runtime - vasp_version = None if runtime_data is None else runtime_data.vasp_version - return { - "vasp_version": vasp_version, - } - - def _dict_from_structure(self) -> dict: - num_ion_steps = None - with suppress(exception.NoData, AttributeError): - positions = self._read_attr("structure", "positions") - if not check.is_none(positions): - num_ion_steps = 1 if positions.ndim == 2 else positions.shape[0] - return { - "num_ionic_steps": num_ion_steps, - } - - def _dict_from_contcar(self) -> dict: - has_selective_dynamics = None - has_lattice_velocities = None - has_ion_velocities = None - with suppress(exception.NoData): - has_selective_dynamics = not check.is_none( - self._read_attr("contcar", "selective_dynamics") - ) - has_lattice_velocities = not check.is_none( - self._read_attr("contcar", "lattice_velocities") - ) - has_ion_velocities = not check.is_none( - self._read_attr("contcar", "ion_velocities") - ) - return { - "has_selective_dynamics": has_selective_dynamics, - "has_lattice_velocities": has_lattice_velocities, - "has_ion_velocities": has_ion_velocities, - } - - def _dict_from_phonon_dispersion(self) -> dict: - phonon_num_qpoints = None - phonon_num_modes = None - with suppress(exception.NoData): - eigenvalues = self._raw_run_info.phonon_dispersion.eigenvalues - phonon_num_qpoints = eigenvalues.shape[0] - phonon_num_modes = eigenvalues.shape[1] - return { - "phonon_num_qpoints": phonon_num_qpoints, - "phonon_num_modes": phonon_num_modes, - } - - -@quantity("run_info") -class RunInfo: - "Contains information about the VASP run." - - def __init__(self, source, quantity_name: str = "run_info"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_run_info: raw.RunInfo) -> "RunInfo": - """Create a RunInfo dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_run_info)) - - def read(self) -> dict: - "Convert the run information to a dictionary." - return merge_default( - self._source, - self._quantity_name, - None, - RunInfoHandler.from_data, - RunInfoHandler.to_dict, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - RunInfoHandler.from_data, - RunInfoHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from contextlib import suppress + +from py4vasp import raw +from py4vasp._calculation import bandgap as bandgap_module +from py4vasp._calculation import exception +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import RunInfo_DB +from py4vasp._raw.data_wrapper import VaspData +from py4vasp._util import check + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + exception.OutdatedVaspVersion, + exception.NoData, + AttributeError, + TypeError, + ValueError, +) + + +class RunInfoHandler: + """Handler for run info. Works with exactly one raw.RunInfo object.""" + + def __init__(self, raw_run_info: raw.RunInfo): + self._raw_run_info = raw_run_info + + @classmethod + def from_data(cls, raw_run_info: raw.RunInfo) -> "RunInfoHandler": + return cls(raw_run_info) + + def to_dict(self) -> dict: + """Convert the run information to a dictionary.""" + return { + **self._dict_from_runtime(), + **self._dict_from_system(), + **self._dict_from_structure(), + **self._dict_additional_collection(), + **self._dict_from_contcar(), + **self._dict_from_phonon_dispersion(), + } + + def to_database(self) -> dict: + """Serialize run info for the database.""" + return RunInfo_DB(**self.to_dict()) + + def _read_attr(self, *keys: str): + data = self._raw_run_info + for key in keys: + data = getattr(data, key, None) + if data is None: + return None + try: + is_none = data.is_none() or data is None + except AttributeError: + is_none = False + return data[:] if not is_none else None + + def _dict_additional_collection(self) -> dict: + fermi_energy = None + with suppress(exception.NoData): + fermi_energy = self._raw_run_info.fermi_energy + if isinstance(fermi_energy, VaspData): + fermi_energy = fermi_energy._data + + is_success = None # TODO implement + + is_collinear = self._is_collinear() + is_noncollinear = self._is_noncollinear() + is_metallic = self._is_metallic() + is_magnetic = None # TODO implement + magnetic_order = None # TODO implement + + grid_coarse_shape = ( + None # TODO implement for FFT grid (currently not written to H5) + ) + grid_fine_shape = ( + None # TODO implement for FFT grid (currently not written to H5) + ) + + return { + "grid_coarse_shape": grid_coarse_shape, + "grid_fine_shape": grid_fine_shape, + "is_success": is_success, + "fermi_energy": fermi_energy, + "is_collinear": is_collinear, + "is_noncollinear": is_noncollinear, + "is_metallic": is_metallic, + "is_magnetic": is_magnetic, + "magnetization_order": magnetic_order, + } + + def _is_collinear(self): + if not check.is_none(self._raw_run_info.len_dos): + return self._raw_run_info.len_dos == 2 + else: + if not check.is_none(self._raw_run_info.band_dispersion_eigenvalues): + return len(self._raw_run_info.band_dispersion_eigenvalues) == 2 + else: + return None + + def _is_noncollinear(self): + if not check.is_none(self._raw_run_info.len_dos): + return self._raw_run_info.len_dos == 4 + else: + if not check.is_none(self._raw_run_info.band_projections): + return len(self._raw_run_info.band_projections) == 4 + else: + return None + + def _is_metallic(self): + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + if check.is_none(self._raw_run_info.bandgap): + return None + gap = bandgap_module.BandgapHandler.from_data(self._raw_run_info.bandgap) + return all(gap._output_gap("fundamental", to_string=False) <= 0.0) + return None + + def _dict_from_system(self) -> dict: + system_tag = None + with suppress(exception.NoData): + system_tag = self._read_attr("system", "system") + return { + "system_tag": system_tag, + } + + def _dict_from_runtime(self) -> dict: + vasp_version = None + with suppress(exception.NoData): + runtime_data = self._raw_run_info.runtime + vasp_version = None if runtime_data is None else runtime_data.vasp_version + return { + "vasp_version": vasp_version, + } + + def _dict_from_structure(self) -> dict: + num_ion_steps = None + with suppress(exception.NoData, AttributeError): + positions = self._read_attr("structure", "positions") + if not check.is_none(positions): + num_ion_steps = 1 if positions.ndim == 2 else positions.shape[0] + return { + "num_ionic_steps": num_ion_steps, + } + + def _dict_from_contcar(self) -> dict: + has_selective_dynamics = None + has_lattice_velocities = None + has_ion_velocities = None + with suppress(exception.NoData): + has_selective_dynamics = not check.is_none( + self._read_attr("contcar", "selective_dynamics") + ) + has_lattice_velocities = not check.is_none( + self._read_attr("contcar", "lattice_velocities") + ) + has_ion_velocities = not check.is_none( + self._read_attr("contcar", "ion_velocities") + ) + return { + "has_selective_dynamics": has_selective_dynamics, + "has_lattice_velocities": has_lattice_velocities, + "has_ion_velocities": has_ion_velocities, + } + + def _dict_from_phonon_dispersion(self) -> dict: + phonon_num_qpoints = None + phonon_num_modes = None + with suppress(exception.NoData): + eigenvalues = self._raw_run_info.phonon_dispersion.eigenvalues + phonon_num_qpoints = eigenvalues.shape[0] + phonon_num_modes = eigenvalues.shape[1] + return { + "phonon_num_qpoints": phonon_num_qpoints, + "phonon_num_modes": phonon_num_modes, + } + + +@quantity("run_info") +class RunInfo: + "Contains information about the VASP run." + + def __init__(self, source, quantity_name: str = "run_info"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_run_info: raw.RunInfo) -> "RunInfo": + """Create a RunInfo dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_run_info)) + + def read(self) -> dict: + "Convert the run information to a dictionary." + return merge_default( + self._source, + self._quantity_name, + None, + RunInfoHandler.from_data, + RunInfoHandler.to_dict, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + RunInfoHandler.from_data, + RunInfoHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/stress.py b/src/py4vasp/_calculation/stress.py index 270109a2c..5b1847c2b 100644 --- a/src/py4vasp/_calculation/stress.py +++ b/src/py4vasp/_calculation/stress.py @@ -1,250 +1,250 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy - -import numpy as np - -from py4vasp import raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, - slice_steps, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw.data_db import Stress_DB -from py4vasp._util import tensor - - -class StressHandler: - """Handler for stress data — performs all data access and transformation logic.""" - - def __init__(self, raw_stress: raw.Stress, steps=None): - self._raw_stress = raw_stress - self._steps = steps - - @classmethod - def from_data(cls, raw_stress: raw.Stress, steps=None) -> "StressHandler": - return cls(raw_stress, steps=steps) - - def __str__(self) -> str: - "Convert the stress to a format similar to the OUTCAR file." - step = self._last_step - structure = StructureHandler.from_data(self._raw_stress.structure, steps=step) - eV_to_kB = 1.602176634e3 / structure.volume() - stress_arr = np.array(self._raw_stress.stress)[step] - stress = tensor.symmetry_reduce(stress_arr) - stress_to_string = lambda s: " ".join(f"{x:11.5f}" for x in s) - return f""" -FORCE on cell =-STRESS in cart. coord. units (eV): -Direction XX YY ZZ XY YZ ZX -------------------------------------------------------------------------------------- -Total {stress_to_string(stress / eV_to_kB)} -in kB {stress_to_string(stress)} -""".strip() - - def to_dict(self) -> dict: - """Read the stress and associated structural information for one or more - selected steps of the trajectory. - - Returns - ------- - dict - Contains the stress for all selected steps and the structural information - to know on which cell the stress acts. - """ - structure = StructureHandler.from_data( - self._raw_stress.structure, steps=self._steps - ) - return { - "stress": slice_steps( - np.array(self._raw_stress.stress), self._steps, default_ndim=2 - ), - "structure": structure.to_dict(), - } - - def to_database(self) -> dict: - """Serialize stress statistics to the database format.""" - stress = np.array(self._raw_stress.stress) - if stress.ndim == 3: - initial_stress_tensor = stress[0] - final_stress_tensor = stress[-1] - else: - initial_stress_tensor = stress - final_stress_tensor = stress - return Stress_DB( - initial_stress_mean=np.trace(initial_stress_tensor) / 3.0, - final_stress_mean=np.trace(final_stress_tensor) / 3.0, - final_stress_tensor=tensor.symmetry_reduce(final_stress_tensor), - ) - - def number_steps(self) -> int: - """Return the number of stress components in the trajectory.""" - n = len(np.array(self._raw_stress.stress)) - return len(range(n)[self._to_slice]) - - @property - def _last_step(self): - if self._steps is None or self._steps == -1: - return -1 - if isinstance(self._steps, slice): - stop = self._steps.stop - return (stop - 1) if stop is not None else -1 - return self._steps - - @property - def _to_slice(self): - if self._steps is None or self._steps == -1: - return slice(-1, None) - if isinstance(self._steps, slice): - return self._steps - return slice(self._steps, self._steps + 1) - - -@quantity("stress") -class Stress: - """The stress describes the force acting on the shape of the unit cell. - - The stress refers to the force applied to the cell per unit area. Specifically, - VASP computes the stress for a given unit cell and relaxing to vanishing stress - determines the predicted ground-state cell. The stress is 3 x 3 matrix; the trace - indicates changes to the volume and the rest of the matrix changes the shape of - the cell. You can impose an external stress with the tag :tag:`PSTRESS`. - - When you relax the system or in a MD simulation, VASP computes and stores the - stress in every iteration. You can use this class to read the stress for specific - steps along the trajectory. - - Examples - -------- - Let us create some example data so that we can illustrate how to use this class. - Of course you can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you access the stress, the result will depend on the steps that you selected - with the [] operator. Without any selection the results from the final step will be - used. - - >>> calculation.stress.number_steps() - 1 - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.stress[:].number_steps() - 4 - - You can also select specific steps or a subset of steps as follows - - >>> calculation.stress[3].number_steps() - 1 - >>> calculation.stress[1:4].number_steps() - 3 - """ - - def __init__(self, source, quantity_name: str = "stress", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_stress: raw.Stress) -> "Stress": - """Create a Stress dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_stress)) - - def __getitem__(self, steps) -> "Stress": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw): - return StressHandler.from_data(raw, steps=self._steps) - - def __str__(self, selection=None) -> str: - "Convert the stress to a format similar to the OUTCAR file." - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - StressHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read() - - def read(self) -> dict: - """Read the stress and associated structural information for one or more - selected steps of the trajectory. - - Returns - ------- - dict - Contains the stress for all selected steps and the structural information - to know on which cell the stress acts. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `read` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. The structure is included to provide the necessary context for - the stress. - - >>> calculation.stress.read() - {'structure': {...}, 'stress': array([[...]])} - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the stress contains an additional dimension for the - different steps. - - >>> calculation.stress[:].read() - {'structure': {...}, 'stress': array([[[...]]])} - - You can also select specific steps or a subset of steps as follows - - >>> calculation.stress[1].read() - {'structure': {...}, 'stress': array([[...]])} - >>> calculation.stress[0:2].read() - {'structure': {...}, 'stress': array([[[...]]])} - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StressHandler.to_dict, - ) - - def number_steps(self) -> int: - """Return the number of stress components in the trajectory.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StressHandler.number_steps, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - StressHandler.from_data, - StressHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy + +import numpy as np + +from py4vasp import raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, + slice_steps, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw.data_db import Stress_DB +from py4vasp._util import tensor + + +class StressHandler: + """Handler for stress data — performs all data access and transformation logic.""" + + def __init__(self, raw_stress: raw.Stress, steps=None): + self._raw_stress = raw_stress + self._steps = steps + + @classmethod + def from_data(cls, raw_stress: raw.Stress, steps=None) -> "StressHandler": + return cls(raw_stress, steps=steps) + + def __str__(self) -> str: + "Convert the stress to a format similar to the OUTCAR file." + step = self._last_step + structure = StructureHandler.from_data(self._raw_stress.structure, steps=step) + eV_to_kB = 1.602176634e3 / structure.volume() + stress_arr = np.array(self._raw_stress.stress)[step] + stress = tensor.symmetry_reduce(stress_arr) + stress_to_string = lambda s: " ".join(f"{x:11.5f}" for x in s) + return f""" +FORCE on cell =-STRESS in cart. coord. units (eV): +Direction XX YY ZZ XY YZ ZX +------------------------------------------------------------------------------------- +Total {stress_to_string(stress / eV_to_kB)} +in kB {stress_to_string(stress)} +""".strip() + + def to_dict(self) -> dict: + """Read the stress and associated structural information for one or more + selected steps of the trajectory. + + Returns + ------- + dict + Contains the stress for all selected steps and the structural information + to know on which cell the stress acts. + """ + structure = StructureHandler.from_data( + self._raw_stress.structure, steps=self._steps + ) + return { + "stress": slice_steps( + np.array(self._raw_stress.stress), self._steps, default_ndim=2 + ), + "structure": structure.to_dict(), + } + + def to_database(self) -> dict: + """Serialize stress statistics to the database format.""" + stress = np.array(self._raw_stress.stress) + if stress.ndim == 3: + initial_stress_tensor = stress[0] + final_stress_tensor = stress[-1] + else: + initial_stress_tensor = stress + final_stress_tensor = stress + return Stress_DB( + initial_stress_mean=np.trace(initial_stress_tensor) / 3.0, + final_stress_mean=np.trace(final_stress_tensor) / 3.0, + final_stress_tensor=tensor.symmetry_reduce(final_stress_tensor), + ) + + def number_steps(self) -> int: + """Return the number of stress components in the trajectory.""" + n = len(np.array(self._raw_stress.stress)) + return len(range(n)[self._to_slice]) + + @property + def _last_step(self): + if self._steps is None or self._steps == -1: + return -1 + if isinstance(self._steps, slice): + stop = self._steps.stop + return (stop - 1) if stop is not None else -1 + return self._steps + + @property + def _to_slice(self): + if self._steps is None or self._steps == -1: + return slice(-1, None) + if isinstance(self._steps, slice): + return self._steps + return slice(self._steps, self._steps + 1) + + +@quantity("stress") +class Stress: + """The stress describes the force acting on the shape of the unit cell. + + The stress refers to the force applied to the cell per unit area. Specifically, + VASP computes the stress for a given unit cell and relaxing to vanishing stress + determines the predicted ground-state cell. The stress is 3 x 3 matrix; the trace + indicates changes to the volume and the rest of the matrix changes the shape of + the cell. You can impose an external stress with the tag :tag:`PSTRESS`. + + When you relax the system or in a MD simulation, VASP computes and stores the + stress in every iteration. You can use this class to read the stress for specific + steps along the trajectory. + + Examples + -------- + Let us create some example data so that we can illustrate how to use this class. + Of course you can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you access the stress, the result will depend on the steps that you selected + with the [] operator. Without any selection the results from the final step will be + used. + + >>> calculation.stress.number_steps() + 1 + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.stress[:].number_steps() + 4 + + You can also select specific steps or a subset of steps as follows + + >>> calculation.stress[3].number_steps() + 1 + >>> calculation.stress[1:4].number_steps() + 3 + """ + + def __init__(self, source, quantity_name: str = "stress", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_stress: raw.Stress) -> "Stress": + """Create a Stress dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_stress)) + + def __getitem__(self, steps) -> "Stress": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw): + return StressHandler.from_data(raw, steps=self._steps) + + def __str__(self, selection=None) -> str: + "Convert the stress to a format similar to the OUTCAR file." + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + StressHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read() + + def read(self) -> dict: + """Read the stress and associated structural information for one or more + selected steps of the trajectory. + + Returns + ------- + dict + Contains the stress for all selected steps and the structural information + to know on which cell the stress acts. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `read` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. The structure is included to provide the necessary context for + the stress. + + >>> calculation.stress.read() + {'structure': {...}, 'stress': array([[...]])} + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the stress contains an additional dimension for the + different steps. + + >>> calculation.stress[:].read() + {'structure': {...}, 'stress': array([[[...]]])} + + You can also select specific steps or a subset of steps as follows + + >>> calculation.stress[1].read() + {'structure': {...}, 'stress': array([[...]])} + >>> calculation.stress[0:2].read() + {'structure': {...}, 'stress': array([[[...]]])} + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StressHandler.to_dict, + ) + + def number_steps(self) -> int: + """Return the number of stress components in the trajectory.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StressHandler.number_steps, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + StressHandler.from_data, + StressHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/structure.py b/src/py4vasp/_calculation/structure.py index ed1efd3a5..f34244780 100644 --- a/src/py4vasp/_calculation/structure.py +++ b/src/py4vasp/_calculation/structure.py @@ -1,1274 +1,1274 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy -from contextlib import suppress -from dataclasses import dataclass -from typing import Union - -import numpy as np - -from py4vasp import exception, raw -from py4vasp._calculation import _stoichiometry -from py4vasp._calculation._stoichiometry import StoichiometryHandler -from py4vasp._calculation.cell import CellHandler -from py4vasp._calculation.dispatch import ( - DataSource, - merge_default, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import Structure_DB -from py4vasp._third_party import view -from py4vasp._util import check, import_, parse - -ase = import_.optional("ase") -ase_io = import_.optional("ase.io") -mdtraj = import_.optional("mdtraj") - -__all__ = ["Structure"] - -_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( - exception.Py4VaspError, - np.linalg.LinAlgError, - AttributeError, - TypeError, - ValueError, - IndexError, -) - - -class StructureHandler: - """Processes structural data from a single raw.Structure object.""" - - A_to_nm = 0.1 - - def __init__(self, raw_structure: raw.Structure, steps=None): - self._raw_structure = raw_structure - self._steps = steps if steps is not None else -1 - self._is_slice = isinstance(self._steps, slice) - if self._is_slice: - self._slice = self._steps - elif self._steps == -1: - self._slice = slice(-1, None) - else: - try: - self._slice = slice(self._steps, self._steps + 1) - except TypeError as error: - raise exception.IncorrectUsage( - f"Steps must be an integer or slice, got {type(self._steps).__name__!r}." - ) from error - - @classmethod - def from_data(cls, raw_structure: raw.Structure, steps=None) -> "StructureHandler": - return cls(raw_structure, steps=steps) - - def to_dict(self, ion_types=None) -> dict: - """Read the structural information into a dictionary. - - The returned dictionary contains the following keys: - - 'lattice_vectors': The lattice vectors of the unit cell. - - 'positions': The positions of the atoms in the unit cell. - - 'elements': The chemical elements of the atoms in the unit cell. - - 'names': The names of the atoms in the unit cell. - - Note that 'elements' and 'names' have the same length as the number of atoms in - the unit cell. - - Parameters - ---------- - ion_types : Sequence - Overwrite the ion types present in the raw data. You can use this to quickly - generate different stoichiometries without modifying the underlying raw data. - - Returns - ------- - dict - Contains the unit cell of the crystal, as well as the position of - all the atoms in units of the lattice vectors and the elements of - the atoms for all selected steps. - """ - return { - "lattice_vectors": self.lattice_vectors(), - "positions": self.positions(), - "elements": self._stoichiometry().elements(ion_types), - "names": self._stoichiometry().names(ion_types), - } - - def __str__(self): - """Generate a string representing the final structure usable as a POSCAR file.""" - return self._create_repr() - - def _repr_html_(self): - format_ = _Format( - begin_table="\n\n\n
", - column_separator="", - row_separator="
", - end_table="
", - newline="
", - ) - return self._create_repr(format_) - - def to_POSCAR(self, ion_types=None) -> str: - """Convert the structure(s) to a POSCAR format.""" - if not self._is_slice: - return self._create_repr(ion_types=ion_types) - else: - message = "Converting multiple structures to a POSCAR is currently not implemented." - raise exception.NotImplemented(message) - - def to_view(self, supercell=None, ion_types=None): - """Generate a 3d representation of the structure(s).""" - make_3d = lambda array: array if array.ndim == 3 else array[np.newaxis] - positions = make_3d(self.positions()) - elements_single_step = self._stoichiometry().elements(ion_types) - elements_all_steps = np.tile(elements_single_step, (len(positions), 1)) - return view.View( - elements=elements_all_steps, - lattice_vectors=make_3d(self.lattice_vectors()), - positions=positions, - supercell=self._parse_supercell(supercell), - ) - - def to_ase(self, supercell=None, ion_types=None): - """Convert the structure to an ASE Atoms object.""" - if self._is_slice: - message = ( - "Converting multiple structures to ASE trajectories is not implemented." - ) - raise exception.NotImplemented(message) - data = self.to_dict(ion_types) - structure = ase.Atoms( - symbols=data["elements"], - cell=data["lattice_vectors"], - scaled_positions=data["positions"], - pbc=True, - ) - num_atoms_prim = len(structure) - if supercell is not None: - try: - structure *= supercell - except (TypeError, IndexError) as err: - error_message = ( - "Generating the supercell failed. Please make sure the requested " - "supercell is either an integer or a list of 3 integers." - ) - raise exception.IncorrectUsage(error_message) from None - num_atoms_super = len(structure) - order = sorted(range(num_atoms_super), key=lambda n: n % num_atoms_prim) - return structure[order] - - def to_mdtraj(self, ion_types=None): - """Convert the trajectory to mdtraj.Trajectory.""" - if not self._is_slice: - message = "Converting a single structure to mdtraj is not implemented." - raise exception.NotImplemented(message) - data = self.to_dict(ion_types) - xyz = data["positions"] @ data["lattice_vectors"] * self.A_to_nm - trajectory = mdtraj.Trajectory(xyz, self._stoichiometry().to_mdtraj(ion_types)) - trajectory.unitcell_vectors = data["lattice_vectors"] * self.A_to_nm - return trajectory - - def to_lammps(self, standard_form=True) -> str: - """Convert the structure to LAMMPS format.""" - if self._is_slice: - message = "Converting multiple structures to LAMMPS is not implemented." - raise exception.NotImplemented(message) - number_ion_types = self._raw_structure.stoichiometry.number_ion_types - cell_string, transformation = self._cell_and_transformation(standard_form) - position_lines = self._position_lines(number_ion_types, transformation) - return f"""\ -Configuration 1: system "{self._stoichiometry()}" - -{self.number_atoms()} atoms -{len(number_ion_types)} atom types - -{cell_string} - -Atoms # atomic - -{position_lines}""" - - def lattice_vectors(self): - """Return the lattice vectors spanning the unit cell.""" - return self._cell().lattice_vectors() - - def positions(self): - """Return the direct coordinates of all ions in the unit cell.""" - return self._raw_structure.positions[self._get_steps()] - - def cartesian_positions(self): - """Convert the positions from direct coordinates to cartesian ones.""" - return self.positions() @ self.lattice_vectors() - - def volume(self): - """Return the volume of the unit cell for the selected steps.""" - return np.abs(np.linalg.det(self.lattice_vectors())) - - def number_atoms(self) -> int: - """Return the total number of atoms in the structure.""" - if self._is_trajectory: - return self._raw_structure.positions.shape[1] - else: - return self._raw_structure.positions.shape[0] - - def number_steps(self) -> int: - """Return the number of structures in the trajectory.""" - if self._is_trajectory: - range_ = range(len(self._raw_structure.positions)) - return len(range_[self._slice]) - else: - return 1 - - def to_database(self) -> dict: - """Return database-ready structure data.""" - # Temporarily use all steps for database - saved_steps = self._steps - saved_is_slice = self._is_slice - saved_slice = self._slice - self._steps = slice(None) - self._is_slice = True - self._slice = slice(None) - - final_lattice, initial_lattice = ([None, None, None] for _ in range(2)) - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - lattices = self.lattice_vectors() - final_lattice = lattices[-1] if lattices.ndim == 3 else lattices - initial_lattice = lattices[0] if lattices.ndim == 3 else lattices - if final_lattice.ndim != 2: - final_lattice = [None, None, None] - if initial_lattice.ndim != 2: - initial_lattice = [None, None, None] - - volume_final, volume_initial = (None for _ in range(2)) - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - volumes = self.volume() - volume_final = ( - volumes[-1] - if not isinstance(volumes, (float, np.float64, np.float32)) - else volumes - ) - volume_initial = ( - volumes[0] - if not isinstance(volumes, (float, np.float64, np.float32)) - else volumes - ) - - lengths_final, angles_final, lengths_initial, angles_initial = ( - None for _ in range(4) - ) - ( - cell_area_2d_final, - cell_area_2d_span_final, - cell_area_2d_initial, - cell_area_2d_span_initial, - ) = (None for _ in range(4)) - dimensionality = 3 - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - dimensionality = self._dimensionality() - - with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): - cell_ = self._cell() - lengths = cell_.lengths() - lengths_final = lengths[-1] if lengths.ndim == 2 else lengths - lengths_initial = lengths[0] if lengths.ndim == 2 else lengths - angles = cell_.angles() - angles_final = angles[-1] if angles.ndim == 2 else angles - angles_initial = angles[0] if angles.ndim == 2 else angles - if dimensionality == 2: - cell_area_2d, cell_area_2d_span = cell_._area_2d() - cell_area_2d_final = ( - cell_area_2d[-1] - if isinstance(cell_area_2d, np.ndarray) - else cell_area_2d - ) - cell_area_2d_initial = ( - cell_area_2d[0] - if isinstance(cell_area_2d, np.ndarray) - else cell_area_2d - ) - cell_area_2d_span_final = ( - cell_area_2d_span[-1] - if isinstance(cell_area_2d_span, list) - else cell_area_2d_span - ) - cell_area_2d_span_initial = ( - cell_area_2d_span[0] - if isinstance(cell_area_2d_span, list) - else cell_area_2d_span - ) - - num_atoms = self.number_atoms() or None - - # Restore steps - self._steps = saved_steps - self._is_slice = saved_is_slice - self._slice = saved_slice - - return Structure_DB( - num_ions=num_atoms, - dimensionality=dimensionality, - final_cell_volume=volume_final, - final_cell_area_2d=cell_area_2d_final, - final_cell_area_2d_span=cell_area_2d_span_final, - final_lattice_vector_1=( - list(final_lattice[0]) if final_lattice[0] is not None else None - ), - final_lattice_vector_2=( - list(final_lattice[1]) if final_lattice[1] is not None else None - ), - final_lattice_vector_3=( - list(final_lattice[2]) if final_lattice[2] is not None else None - ), - final_lattice_vector_1_length=( - lengths_final[0] if lengths_final is not None else None - ), - final_lattice_vector_2_length=( - lengths_final[1] if lengths_final is not None else None - ), - final_lattice_vector_3_length=( - lengths_final[2] if lengths_final is not None else None - ), - final_angle_alpha=(angles_final[0] if angles_final is not None else None), - final_angle_beta=(angles_final[1] if angles_final is not None else None), - final_angle_gamma=(angles_final[2] if angles_final is not None else None), - initial_cell_volume=volume_initial, - initial_cell_area_2d=cell_area_2d_initial, - initial_cell_area_2d_span=cell_area_2d_span_initial, - initial_lattice_vector_1=( - list(initial_lattice[0]) if initial_lattice[0] is not None else None - ), - initial_lattice_vector_2=( - list(initial_lattice[1]) if initial_lattice[1] is not None else None - ), - initial_lattice_vector_3=( - list(initial_lattice[2]) if initial_lattice[2] is not None else None - ), - initial_lattice_vector_1_length=( - lengths_initial[0] if lengths_initial is not None else None - ), - initial_lattice_vector_2_length=( - lengths_initial[1] if lengths_initial is not None else None - ), - initial_lattice_vector_3_length=( - lengths_initial[2] if lengths_initial is not None else None - ), - initial_angle_alpha=( - angles_initial[0] if angles_initial is not None else None - ), - initial_angle_beta=( - angles_initial[1] if angles_initial is not None else None - ), - initial_angle_gamma=( - angles_initial[2] if angles_initial is not None else None - ), - ) - - def _dimensionality(self) -> Union[int, np.ndarray]: - """Heuristic check for dimensionality of system.""" - if not check.is_none(self._raw_structure.idipol): - if self._raw_structure.idipol < 1: - return 3 - elif self._raw_structure.idipol in [1, 2, 3]: - return 2 - elif self._raw_structure.idipol == 4: - return 0 - cell_ = self._cell() - if bool(np.all(np.array(cell_.is_suspected_2d_system))): - return 2 - return 3 - - def _stoichiometry(self) -> StoichiometryHandler: - return StoichiometryHandler.from_data(self._raw_structure.stoichiometry) - - def _cell(self) -> CellHandler: - return CellHandler.from_data(self._raw_structure.cell, steps=self._steps) - - def _get_steps(self): - return self._steps if self._is_trajectory else () - - def _get_last_step(self): - return self._last_step_in_slice if self._is_trajectory else () - - @property - def _last_step_in_slice(self): - return (self._slice.stop or 0) - 1 - - def _step_string(self): - if self._is_slice: - range_ = range(len(self._raw_structure.positions))[self._steps] - return f" from step {range_.start + 1} to {range_.stop + 1}" - elif self._steps == -1: - return "" - else: - return f" (step {self._steps + 1})" - - def _create_repr(self, format_=None, ion_types=None): - if format_ is None: - format_ = _Format() - step = self._get_last_step() - stoichiometry = self._stoichiometry() - lines = ( - format_.comment_line(stoichiometry, self._step_string(), ion_types), - format_.scaling_factor(self._cell().scale()), - format_.vectors_to_table(self._raw_structure.cell.lattice_vectors[step]), - format_.ion_list(stoichiometry, ion_types), - format_.coordinate_system(), - format_.vectors_to_table(self._raw_structure.positions[step]), - ) - return "\n".join(lines) - - def _parse_supercell(self, supercell): - if supercell is None: - return np.ones(3, np.int_) - try: - integer_supercell = np.round(supercell).astype(np.int_) - except TypeError as error: - message = ( - f"Could not convert supercell='{supercell}' to an integer numpy array." - ) - raise exception.IncorrectUsage(message) from None - if not np.allclose(supercell, integer_supercell): - message = f"supercell='{supercell}' contains noninteger values." - raise exception.IncorrectUsage(message) - if np.isscalar(integer_supercell): - return np.full(3, integer_supercell) - if integer_supercell.shape == (3,): - return integer_supercell - message = ( - f"supercell='{supercell}' is not a scalar or a three component vector." - ) - raise exception.IncorrectUsage(message) - - def _cell_and_transformation(self, standard_form): - if standard_form: - cell_obj = ase.cell.Cell(self.lattice_vectors()) - cell_obj, transformation = cell_obj.standard_form() - cell_string = f"""\ -0.0 {self._format_number(cell_obj[0,0])} xlo xhi -0.0 {self._format_number(cell_obj[1,1])} ylo yhi -0.0 {self._format_number(cell_obj[2,2])} zlo zhi -{self._format_number((cell_obj[1,0], cell_obj[2,0], cell_obj[2,1]))} xy xz yz""" - else: - lattice_vectors = self.lattice_vectors() - cell_string = f"""\ -{self._format_number(lattice_vectors[0])} avec -{self._format_number(lattice_vectors[1])} bvec -{self._format_number(lattice_vectors[2])} cvec -0.0 0.0 0.0 abc origin""" - transformation = np.eye(3) - return cell_string, transformation - - def _position_lines(self, number_ion_types, transformation): - positions = self.cartesian_positions() @ transformation.T - ion_type_labels = [ - str(ion_type + 1) - for ion_type, number in enumerate(number_ion_types) - for _ in range(number) - ] - return "\n".join( - f"{i + 1} {ion_type_labels[i]} {self._format_number(position)}" - for i, position in enumerate(positions) - ) - - def _format_number(self, number): - number = np.atleast_1d(number) - return " ".join(f"{x:24.16E}" for x in number) - - @property - def _is_trajectory(self): - return self._raw_structure.positions.ndim == 3 - - -@dataclass -class _Format: - begin_table: str = "" - column_separator: str = " " - row_separator: str = "\n" - end_table: str = "" - newline: str = "" - - def comment_line(self, stoichiometry, step_string, ion_types): - return f"{stoichiometry.to_string(ion_types)}{step_string}{self.newline}" - - def scaling_factor(self, scale): - return f"{self._element_to_string(scale)}{self.newline}".lstrip() - - def ion_list(self, stoichiometry, ion_types): - return f"{stoichiometry.to_POSCAR(self.newline, ion_types)}{self.newline}" - - def coordinate_system(self): - return f"Direct{self.newline}" - - def vectors_to_table(self, vectors): - rows = (self._vector_to_row(vector) for vector in vectors) - return f"{self.begin_table}{self.row_separator.join(rows)}{self.end_table}" - - def _vector_to_row(self, vector): - elements = (self._element_to_string(element) for element in vector) - return self.column_separator.join(elements) - - def _element_to_string(self, element): - return f"{element:21.16f}" - - -@quantity("structure") -class Structure(view.Mixin): - """The structure contains the unit cell and the position of all ions within. - - The crystal structure is the specific arrangement of ions in a three-dimensional - repeating pattern. This spatial arrangement is characterized by the unit cell and - the relative position of the ions. The unit cell is repeated periodically in three - dimensions to form the crystal. The combination of unit cell and ion positions - determines the symmetry of the crystal. This symmetry helps understanding the - material properties because some symmetries do not allow for the presence of some - properties, e.g., you cannot observe a ferroelectric :data:`~py4vasp.calculation.polarization` - in a system with inversion symmetry. Therefore relaxing the crystal structure with - VASP is an important first step in analyzing materials properties. - - When you run a relaxation or MD simulation, this class allows to access all - individual steps of the trajectory. Typically, you would study the converged - structure after an ionic relaxation or to visualize the changes of the structure - along the simulation. Moreover, you could take snapshots along the trajectory - and further process them by computing more properties. - - Examples - -------- - Let us create some example data so that we can illustrate how to use this class. - Of course you can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you access the structure, the result will depend on the steps that you selected - with the [] operator. Without any selection the results from the final step will be - used. - - >>> calculation.structure.number_steps() - 1 - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.structure[:].number_steps() - 4 - - You can also select specific {step}s or a subset of {step}s as follows - - >>> calculation.structure[3].number_steps() - 1 - >>> calculation.structure[1:4].number_steps() - 3 - """ - - A_to_nm = 0.1 - "Converting Å to nm used for mdtraj trajectories." - - def __init__(self, source, quantity_name="structure", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_structure) -> "Structure": - return cls(source=DataSource(raw_structure)) - - def __repr__(self): - return f"{type(self).__name__}.from_path({self._path!r})" - - def _handler_factory(self, raw): - return StructureHandler.from_data(raw, steps=self._steps) - - def __getitem__(self, steps) -> "Structure": - with self._source.access(self._quantity_name) as raw_structure: - is_trajectory = raw_structure.positions.ndim == 3 - if not is_trajectory: - message = ( - "The structure is not a Trajectory so accessing individual " - "elements is not allowed." - ) - raise exception.IncorrectUsage(message) - new = copy.copy(self) - new._steps = steps - return new - - @classmethod - def from_POSCAR(cls, poscar, *, elements=None): - """Generate a structure from string in POSCAR format. - - The POSCAR format is the standard format to represent crystal structures in - VASP. This method allows to create a structure from a POSCAR string. - To read more about the POSCAR format, please refer to the `VASP manual `_. - - Parameters - ---------- - elements : list[str] - Name of the elements in the order they appear in the POSCAR file. If the - elements are specified in the POSCAR file, this argument is optional and - if set it will overwrite the choice in the POSCAR file. Old POSCAR files - do not specify the name of the elements; in that case this argument is - required. - - Examples - -------- - We can create a GaAs structure from a POSCAR string as follows - - >>> poscar = '''\\ - ... GaAs - ... 5.65325 - ... 0.0 0.5 0.5 - ... 0.5 0.0 0.5 - ... 0.5 0.5 0.0 - ... 1 1 - ... fractional - ... 0.0 0.0 0.0 - ... 0.25 0.25 0.25''' - >>> structure = py4vasp.calculation.structure.from_POSCAR(poscar, elements=['Ga', 'As']) - >>> print(structure.to_POSCAR()) - GaAs - 5.6532... - 0.0... 0.5... 0.5... - 0.5... 0.0... 0.5... - 0.5... 0.5... 0.0... - Ga As - 1 1 - Direct - 0.00... 0.00... 0.00... - 0.25... 0.25... 0.25... - """ - poscar = _replace_or_set_elements(str(poscar), elements) - poscar = parse.POSCAR(poscar) - return cls.from_data(poscar.structure) - - @classmethod - def from_ase(cls, structure): - """Generate a structure from the ase Atoms class.""" - structure = raw.Structure( - stoichiometry=_stoichiometry.raw_stoichiometry_from_ase(structure), - cell=_cell_from_ase(structure), - positions=structure.get_scaled_positions()[np.newaxis], - ) - return cls.from_data(structure) - - def __str__(self, selection=None): - "Generate a string representing the final structure usable as a POSCAR file." - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - StructureHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def _repr_html_(self): - return merge_strings( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler._repr_html_, - ) - - def read(self, ion_types=None): - """Read the structural information into a dictionary. - - The returned dictionary contains the following keys: - - 'lattice_vectors': The lattice vectors of the unit cell. - - 'positions': The positions of the atoms in the unit cell. - - 'elements': The chemical elements of the atoms in the unit cell. - - 'names': The names of the atoms in the unit cell. - - Note that 'elements' and 'names' have the same length as the number of atoms in - the unit cell. - - Parameters - ---------- - ion_types : Sequence - Overwrite the ion types present in the raw data. You can use this to quickly - generate different stoichiometries without modifying the underlying raw data. - - Returns - ------- - dict - Contains the unit cell of the crystal, as well as the position of - all the atoms in units of the lattice vectors and the elements of - the atoms for all selected steps. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `read` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.read() - {'lattice_vectors': array([[...]]), 'positions': array([[...]]), - 'elements': [...], 'names': [...]} - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the lattice vectors and positions contain an additional - dimension for the different steps. - - >>> calculation.structure[:].read() - {'lattice_vectors': array([[[...]]]), 'positions': array([[[...]]]), - 'elements': [...], 'names': [...]} - - You can also select specific steps or a subset of steps as follows - - >>> calculation.structure[1].read() - {'lattice_vectors': array([[...]]), 'positions': array([[...]]), - 'elements': [...], 'names': [...]} - >>> calculation.structure[0:2].read() - {'lattice_vectors': array([[[...]]]), 'positions': array([[[...]]]), - 'elements': [...], 'names': [...]} - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.to_dict, - ion_types, - ) - - def to_dict(self, ion_types=None): - """Convenient alias for :py:meth:`read`. Please read the documentation there.""" - return self.read(ion_types=ion_types) - - def to_view(self, supercell=None, ion_types=None): - """Generate a 3d representation of the structure(s). - - This method uses the `View` class to create a 3d visualization of the atomic - structure(s) in the unit cell. - - Parameters - ---------- - supercell : int or np.ndarray - If present the structure is replicated the specified number of times - along each direction. - ion_types : Sequence - Overwrite the ion types present in the raw data. You can use this to quickly - generate different stoichiometries without modifying the underlying raw data. - - Returns - ------- - View - Visualize the structure(s) as a 3d figure. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_view` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.to_view() - View(elements=array([[...]], dtype=...), lattice_vectors=array([[[...]]]), - positions=array([[[...]]]), ...) - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the lattice vectors and positions contain an additional - dimension for the different steps. - - >>> calculation.structure[:].to_view() - View(elements=array([[...], ..., [...]], dtype=...), lattice_vectors=array([[[...]], ..., [[...]]]), - positions=array([[[...]], ..., [[...]]]), ...) - - You can also select specific steps or a subset of steps as follows - - >>> calculation.structure[1].to_view() - View(elements=array([[...]], dtype=...), lattice_vectors=array([[[...]]]), - positions=array([[[...]]]), ...) - >>> calculation.structure[0:2].to_view() - View(elements=array([[...], [...]], dtype=...), lattice_vectors=array([[[...]], [[...]]]), - positions=array([[[...]], [[...]]]), ...) - - You may also replicate the structure by specifying a supercell. - - >>> calculation.structure.to_view(supercell=2) - View(..., supercell=array([2, 2, 2]), ...) - - The supercell size can also be different for the different directions. - - >>> calculation.structure.to_view(supercell=[2,3,1]) - View(..., supercell=array([2, 3, 1]), ...) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.to_view, - supercell, - ion_types, - ) - - def to_ase(self, supercell=None, ion_types=None): - """Convert the structure to an ASE Atoms object. - - ASE (the Atomic Simulation Environment) is a popular Python package for atomistic - simulations. This method converts the VASP structure to an ASE Atoms object, - which can be used for further analysis and visualization. - - Parameters - ---------- - supercell : int or np.ndarray - If present the structure is replicated the specified number of times - along each direction. - ion_types : Sequence - Overwrite the ion types present in the raw data. You can use this to quickly - generate different stoichiometries without modifying the underlying raw data. - - Returns - ------- - Atoms - Structural information for ASE package. Read more about ASE `here `_. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_ase` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.to_ase() - Atoms(symbols='...', pbc=True, cell=[[...]]) - - You can also select specific steps as follows - - >>> calculation.structure[1].to_ase() - Atoms(symbols='...', pbc=True, cell=[[...]]) - - Notice that converting multiple steps to ASE trajectories is not implemented. - - You may also replicate the structure by specifying a supercell. If you compare - the cell size with the previous example, you will see that it is doubled in all - directions. - - >>> calculation.structure.to_ase(supercell=2) - Atoms(symbols='...', pbc=True, cell=[[...]]) - - The supercell size can also be different for the different directions. The three - lattice vectors will be scaled accordingly. - - >>> calculation.structure.to_ase(supercell=[2,3,1]) - Atoms(symbols='...', pbc=True, cell=[[...]]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.to_ase, - supercell, - ion_types, - ) - - def to_mdtraj(self, ion_types=None): - """Convert the trajectory to mdtraj.Trajectory - - mdtraj is a popular Python package to analyze molecular dynamics trajectories. - This method converts the VASP structure trajectory to an mdtraj.Trajectory - object, which can be used for further analysis and visualization. - - Parameters - ---------- - ion_types : Sequence - Overwrite the ion types present in the raw data. You can use this to quickly - generate different stoichiometries without modifying the underlying raw data. - - Returns - ------- - mdtraj.Trajectory - The mdtraj package offers many functionalities to analyze a MD - trajectory. By converting the VASP data to their format, we facilitate - using all functions of that package. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - To convert the whole trajectory (all steps), you don't specify the array boundaries. - - >>> calculation.structure[:].to_mdtraj() - - - You can also select a subset of steps as follows - - >>> calculation.structure[0:2].to_mdtraj() - - - You cannot convert a single structure to mdtraj.Trajectory. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.to_mdtraj, - ion_types, - ) - - def to_POSCAR(self, ion_types=None): - """Convert the structure(s) to a POSCAR format. - - Use this method to generate a string in POSCAR format representing the - structure(s). You can use this string to write a POSCAR file for VASP. This - can be useful if you want to use the relaxed structure from a VASP calculation - or a snapshot from an MD simulation as input for a new VASP calculation. - - Parameters - ---------- - ion_types : Sequence - Overwrite the ion types present in the raw data. You can use this to quickly - generate different stoichiometries without modifying the underlying raw data. - - Returns - ------- - str - Returns the POSCAR of the selected steps. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_POSCAR` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> poscar = calculation.structure.to_POSCAR() - >>> assert poscar == str(calculation.structure) - - You can also select specific steps as follows - - >>> poscar = calculation.structure[1].to_POSCAR() - >>> assert poscar == str(calculation.structure[1]) - - Notice that converting multiple steps to POSCAR format is not implemented. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.to_POSCAR, - ion_types, - ) - - def to_lammps(self, standard_form=True): - """Convert the structure to LAMMPS format. - - LAMMPS is a popular molecular dynamics simulation software. This method - converts the structure to a string in LAMMPS format, which can be used as - input for LAMMPS simulations. - - Parameters - ---------- - standard_form : bool - Determines whether the structure is standardize, i.e., the lattice vectors - are a triagonal matrix. - - Returns - ------- - str - Returns a string describing the structure for LAMMPS - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_lammps` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> print(calculation.structure.to_lammps()) - Configuration 1: system "..." - ... atoms - ... atom types - ... xlo xhi - ... ylo yhi - ... zlo zhi - ... xy xz yz - Atoms # atomic - 1 1 ... - - You can also select specific steps as follows - - >>> print(calculation.structure[1].to_lammps()) - Configuration 1: system "..." - ... atoms - ... atom types - ... xlo xhi - ... ylo yhi - ... zlo zhi - ... xy xz yz - Atoms # atomic - 1 1 ... - - Notice that converting multiple steps to LAMMPS format is not implemented. - - LAMMPS requires either a standard form of the unit cell or the transformation - from the original cell to the standard form. By default, the standard form is - used. You can disable this behavior as follows - - >>> print(calculation.structure.to_lammps(standard_form=False)) - Configuration 1: system "..." - ... atoms - ... atom types - ... avec - ... bvec - ... cvec - ... abc origin - Atoms # atomic - 1 1 ... - - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.to_lammps, - standard_form, - ) - - def lattice_vectors(self): - """Return the lattice vectors spanning the unit cell - - Returns - ------- - np.ndarray - Lattice vectors of the unit cell in Å. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `lattice_vectors` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.lattice_vectors() - array([[...], [...], [...]]) - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.structure[:].lattice_vectors() - array([[[...]], [[...]], [[...]], [[...]]]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.lattice_vectors, - ) - - def positions(self): - """Return the direct coordinates of all ions in the unit cell. - - Direct or fractional coordinates measure the position of the ions in terms of - the lattice vectors. Hence they are dimensionless quantities. - - Returns - ------- - np.ndarray - Positions of all ions in terms of the lattice vectors. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `positions` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.positions() - array([[...]]) - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.structure[:].positions() - array([[[...]]]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.positions, - ) - - def cartesian_positions(self): - """Convert the positions from direct coordinates to cartesian ones. - - Returns - ------- - np.ndarray - Position of all atoms in cartesian coordinates in Å. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `cartesian_positions` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.cartesian_positions() - array([[...]]) - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.structure[:].cartesian_positions() - array([[[...]]]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.cartesian_positions, - ) - - def volume(self): - """Return the volume of the unit cell for the selected steps. - - Returns - ------- - float or np.ndarray - The volume(s) of the selected step(s) in ų. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this method. - You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `volume` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.structure.volume() - np.float... - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.structure[:].volume() - array([...]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.volume, - ) - - def number_atoms(self): - """Return the total number of atoms in the structure.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.number_atoms, - ) - - def number_steps(self): - """Return the number of structures in the trajectory.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - StructureHandler.number_steps, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - StructureHandler.from_data, - StructureHandler.to_database, - ) - - -def _cell_from_ase(structure): - lattice_vectors = np.array([structure.get_cell()]) - return raw.Cell(lattice_vectors, scale=raw.VaspData(1.0)) - - -def _replace_or_set_elements(poscar, elements): - line_with_elements = 5 - elements = "" if not elements else " ".join(elements) - lines = poscar.split("\n") - if _elements_not_in_poscar(lines[line_with_elements]): - _raise_error_if_elements_not_set(elements) - lines.insert(line_with_elements, elements) - elif elements: - lines[line_with_elements] = elements - return "\n".join(lines) - - -def _elements_not_in_poscar(elements): - elements = elements.split() - return any(element.isdecimal() for element in elements) - - -def _raise_error_if_elements_not_set(elements): - if not elements: - message = """The POSCAR file does not specify the elements needed to create a - Structure. Please pass `elements=[...]` to the `from_POSCAR` routine where - ... are the elements in the same order as in the POSCAR.""" - raise exception.IncorrectUsage(message) - - -class Mixin: - @property - def _structure(self): - return Structure.from_data(self._raw_data.structure) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy +from contextlib import suppress +from dataclasses import dataclass +from typing import Union + +import numpy as np + +from py4vasp import exception, raw +from py4vasp._calculation import _stoichiometry +from py4vasp._calculation._stoichiometry import StoichiometryHandler +from py4vasp._calculation.cell import CellHandler +from py4vasp._calculation.dispatch import ( + DataSource, + merge_default, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import Structure_DB +from py4vasp._third_party import view +from py4vasp._util import check, import_, parse + +ase = import_.optional("ase") +ase_io = import_.optional("ase.io") +mdtraj = import_.optional("mdtraj") + +__all__ = ["Structure"] + +_TO_DATABASE_SUPPRESSED_EXCEPTIONS = ( + exception.Py4VaspError, + np.linalg.LinAlgError, + AttributeError, + TypeError, + ValueError, + IndexError, +) + + +class StructureHandler: + """Processes structural data from a single raw.Structure object.""" + + A_to_nm = 0.1 + + def __init__(self, raw_structure: raw.Structure, steps=None): + self._raw_structure = raw_structure + self._steps = steps if steps is not None else -1 + self._is_slice = isinstance(self._steps, slice) + if self._is_slice: + self._slice = self._steps + elif self._steps == -1: + self._slice = slice(-1, None) + else: + try: + self._slice = slice(self._steps, self._steps + 1) + except TypeError as error: + raise exception.IncorrectUsage( + f"Steps must be an integer or slice, got {type(self._steps).__name__!r}." + ) from error + + @classmethod + def from_data(cls, raw_structure: raw.Structure, steps=None) -> "StructureHandler": + return cls(raw_structure, steps=steps) + + def to_dict(self, ion_types=None) -> dict: + """Read the structural information into a dictionary. + + The returned dictionary contains the following keys: + - 'lattice_vectors': The lattice vectors of the unit cell. + - 'positions': The positions of the atoms in the unit cell. + - 'elements': The chemical elements of the atoms in the unit cell. + - 'names': The names of the atoms in the unit cell. + + Note that 'elements' and 'names' have the same length as the number of atoms in + the unit cell. + + Parameters + ---------- + ion_types : Sequence + Overwrite the ion types present in the raw data. You can use this to quickly + generate different stoichiometries without modifying the underlying raw data. + + Returns + ------- + dict + Contains the unit cell of the crystal, as well as the position of + all the atoms in units of the lattice vectors and the elements of + the atoms for all selected steps. + """ + return { + "lattice_vectors": self.lattice_vectors(), + "positions": self.positions(), + "elements": self._stoichiometry().elements(ion_types), + "names": self._stoichiometry().names(ion_types), + } + + def __str__(self): + """Generate a string representing the final structure usable as a POSCAR file.""" + return self._create_repr() + + def _repr_html_(self): + format_ = _Format( + begin_table="\n\n\n
", + column_separator="", + row_separator="
", + end_table="
", + newline="
", + ) + return self._create_repr(format_) + + def to_POSCAR(self, ion_types=None) -> str: + """Convert the structure(s) to a POSCAR format.""" + if not self._is_slice: + return self._create_repr(ion_types=ion_types) + else: + message = "Converting multiple structures to a POSCAR is currently not implemented." + raise exception.NotImplemented(message) + + def to_view(self, supercell=None, ion_types=None): + """Generate a 3d representation of the structure(s).""" + make_3d = lambda array: array if array.ndim == 3 else array[np.newaxis] + positions = make_3d(self.positions()) + elements_single_step = self._stoichiometry().elements(ion_types) + elements_all_steps = np.tile(elements_single_step, (len(positions), 1)) + return view.View( + elements=elements_all_steps, + lattice_vectors=make_3d(self.lattice_vectors()), + positions=positions, + supercell=self._parse_supercell(supercell), + ) + + def to_ase(self, supercell=None, ion_types=None): + """Convert the structure to an ASE Atoms object.""" + if self._is_slice: + message = ( + "Converting multiple structures to ASE trajectories is not implemented." + ) + raise exception.NotImplemented(message) + data = self.to_dict(ion_types) + structure = ase.Atoms( + symbols=data["elements"], + cell=data["lattice_vectors"], + scaled_positions=data["positions"], + pbc=True, + ) + num_atoms_prim = len(structure) + if supercell is not None: + try: + structure *= supercell + except (TypeError, IndexError) as err: + error_message = ( + "Generating the supercell failed. Please make sure the requested " + "supercell is either an integer or a list of 3 integers." + ) + raise exception.IncorrectUsage(error_message) from None + num_atoms_super = len(structure) + order = sorted(range(num_atoms_super), key=lambda n: n % num_atoms_prim) + return structure[order] + + def to_mdtraj(self, ion_types=None): + """Convert the trajectory to mdtraj.Trajectory.""" + if not self._is_slice: + message = "Converting a single structure to mdtraj is not implemented." + raise exception.NotImplemented(message) + data = self.to_dict(ion_types) + xyz = data["positions"] @ data["lattice_vectors"] * self.A_to_nm + trajectory = mdtraj.Trajectory(xyz, self._stoichiometry().to_mdtraj(ion_types)) + trajectory.unitcell_vectors = data["lattice_vectors"] * self.A_to_nm + return trajectory + + def to_lammps(self, standard_form=True) -> str: + """Convert the structure to LAMMPS format.""" + if self._is_slice: + message = "Converting multiple structures to LAMMPS is not implemented." + raise exception.NotImplemented(message) + number_ion_types = self._raw_structure.stoichiometry.number_ion_types + cell_string, transformation = self._cell_and_transformation(standard_form) + position_lines = self._position_lines(number_ion_types, transformation) + return f"""\ +Configuration 1: system "{self._stoichiometry()}" + +{self.number_atoms()} atoms +{len(number_ion_types)} atom types + +{cell_string} + +Atoms # atomic + +{position_lines}""" + + def lattice_vectors(self): + """Return the lattice vectors spanning the unit cell.""" + return self._cell().lattice_vectors() + + def positions(self): + """Return the direct coordinates of all ions in the unit cell.""" + return self._raw_structure.positions[self._get_steps()] + + def cartesian_positions(self): + """Convert the positions from direct coordinates to cartesian ones.""" + return self.positions() @ self.lattice_vectors() + + def volume(self): + """Return the volume of the unit cell for the selected steps.""" + return np.abs(np.linalg.det(self.lattice_vectors())) + + def number_atoms(self) -> int: + """Return the total number of atoms in the structure.""" + if self._is_trajectory: + return self._raw_structure.positions.shape[1] + else: + return self._raw_structure.positions.shape[0] + + def number_steps(self) -> int: + """Return the number of structures in the trajectory.""" + if self._is_trajectory: + range_ = range(len(self._raw_structure.positions)) + return len(range_[self._slice]) + else: + return 1 + + def to_database(self) -> dict: + """Return database-ready structure data.""" + # Temporarily use all steps for database + saved_steps = self._steps + saved_is_slice = self._is_slice + saved_slice = self._slice + self._steps = slice(None) + self._is_slice = True + self._slice = slice(None) + + final_lattice, initial_lattice = ([None, None, None] for _ in range(2)) + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + lattices = self.lattice_vectors() + final_lattice = lattices[-1] if lattices.ndim == 3 else lattices + initial_lattice = lattices[0] if lattices.ndim == 3 else lattices + if final_lattice.ndim != 2: + final_lattice = [None, None, None] + if initial_lattice.ndim != 2: + initial_lattice = [None, None, None] + + volume_final, volume_initial = (None for _ in range(2)) + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + volumes = self.volume() + volume_final = ( + volumes[-1] + if not isinstance(volumes, (float, np.float64, np.float32)) + else volumes + ) + volume_initial = ( + volumes[0] + if not isinstance(volumes, (float, np.float64, np.float32)) + else volumes + ) + + lengths_final, angles_final, lengths_initial, angles_initial = ( + None for _ in range(4) + ) + ( + cell_area_2d_final, + cell_area_2d_span_final, + cell_area_2d_initial, + cell_area_2d_span_initial, + ) = (None for _ in range(4)) + dimensionality = 3 + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + dimensionality = self._dimensionality() + + with suppress(*_TO_DATABASE_SUPPRESSED_EXCEPTIONS): + cell_ = self._cell() + lengths = cell_.lengths() + lengths_final = lengths[-1] if lengths.ndim == 2 else lengths + lengths_initial = lengths[0] if lengths.ndim == 2 else lengths + angles = cell_.angles() + angles_final = angles[-1] if angles.ndim == 2 else angles + angles_initial = angles[0] if angles.ndim == 2 else angles + if dimensionality == 2: + cell_area_2d, cell_area_2d_span = cell_._area_2d() + cell_area_2d_final = ( + cell_area_2d[-1] + if isinstance(cell_area_2d, np.ndarray) + else cell_area_2d + ) + cell_area_2d_initial = ( + cell_area_2d[0] + if isinstance(cell_area_2d, np.ndarray) + else cell_area_2d + ) + cell_area_2d_span_final = ( + cell_area_2d_span[-1] + if isinstance(cell_area_2d_span, list) + else cell_area_2d_span + ) + cell_area_2d_span_initial = ( + cell_area_2d_span[0] + if isinstance(cell_area_2d_span, list) + else cell_area_2d_span + ) + + num_atoms = self.number_atoms() or None + + # Restore steps + self._steps = saved_steps + self._is_slice = saved_is_slice + self._slice = saved_slice + + return Structure_DB( + num_ions=num_atoms, + dimensionality=dimensionality, + final_cell_volume=volume_final, + final_cell_area_2d=cell_area_2d_final, + final_cell_area_2d_span=cell_area_2d_span_final, + final_lattice_vector_1=( + list(final_lattice[0]) if final_lattice[0] is not None else None + ), + final_lattice_vector_2=( + list(final_lattice[1]) if final_lattice[1] is not None else None + ), + final_lattice_vector_3=( + list(final_lattice[2]) if final_lattice[2] is not None else None + ), + final_lattice_vector_1_length=( + lengths_final[0] if lengths_final is not None else None + ), + final_lattice_vector_2_length=( + lengths_final[1] if lengths_final is not None else None + ), + final_lattice_vector_3_length=( + lengths_final[2] if lengths_final is not None else None + ), + final_angle_alpha=(angles_final[0] if angles_final is not None else None), + final_angle_beta=(angles_final[1] if angles_final is not None else None), + final_angle_gamma=(angles_final[2] if angles_final is not None else None), + initial_cell_volume=volume_initial, + initial_cell_area_2d=cell_area_2d_initial, + initial_cell_area_2d_span=cell_area_2d_span_initial, + initial_lattice_vector_1=( + list(initial_lattice[0]) if initial_lattice[0] is not None else None + ), + initial_lattice_vector_2=( + list(initial_lattice[1]) if initial_lattice[1] is not None else None + ), + initial_lattice_vector_3=( + list(initial_lattice[2]) if initial_lattice[2] is not None else None + ), + initial_lattice_vector_1_length=( + lengths_initial[0] if lengths_initial is not None else None + ), + initial_lattice_vector_2_length=( + lengths_initial[1] if lengths_initial is not None else None + ), + initial_lattice_vector_3_length=( + lengths_initial[2] if lengths_initial is not None else None + ), + initial_angle_alpha=( + angles_initial[0] if angles_initial is not None else None + ), + initial_angle_beta=( + angles_initial[1] if angles_initial is not None else None + ), + initial_angle_gamma=( + angles_initial[2] if angles_initial is not None else None + ), + ) + + def _dimensionality(self) -> Union[int, np.ndarray]: + """Heuristic check for dimensionality of system.""" + if not check.is_none(self._raw_structure.idipol): + if self._raw_structure.idipol < 1: + return 3 + elif self._raw_structure.idipol in [1, 2, 3]: + return 2 + elif self._raw_structure.idipol == 4: + return 0 + cell_ = self._cell() + if bool(np.all(np.array(cell_.is_suspected_2d_system))): + return 2 + return 3 + + def _stoichiometry(self) -> StoichiometryHandler: + return StoichiometryHandler.from_data(self._raw_structure.stoichiometry) + + def _cell(self) -> CellHandler: + return CellHandler.from_data(self._raw_structure.cell, steps=self._steps) + + def _get_steps(self): + return self._steps if self._is_trajectory else () + + def _get_last_step(self): + return self._last_step_in_slice if self._is_trajectory else () + + @property + def _last_step_in_slice(self): + return (self._slice.stop or 0) - 1 + + def _step_string(self): + if self._is_slice: + range_ = range(len(self._raw_structure.positions))[self._steps] + return f" from step {range_.start + 1} to {range_.stop + 1}" + elif self._steps == -1: + return "" + else: + return f" (step {self._steps + 1})" + + def _create_repr(self, format_=None, ion_types=None): + if format_ is None: + format_ = _Format() + step = self._get_last_step() + stoichiometry = self._stoichiometry() + lines = ( + format_.comment_line(stoichiometry, self._step_string(), ion_types), + format_.scaling_factor(self._cell().scale()), + format_.vectors_to_table(self._raw_structure.cell.lattice_vectors[step]), + format_.ion_list(stoichiometry, ion_types), + format_.coordinate_system(), + format_.vectors_to_table(self._raw_structure.positions[step]), + ) + return "\n".join(lines) + + def _parse_supercell(self, supercell): + if supercell is None: + return np.ones(3, np.int_) + try: + integer_supercell = np.round(supercell).astype(np.int_) + except TypeError as error: + message = ( + f"Could not convert supercell='{supercell}' to an integer numpy array." + ) + raise exception.IncorrectUsage(message) from None + if not np.allclose(supercell, integer_supercell): + message = f"supercell='{supercell}' contains noninteger values." + raise exception.IncorrectUsage(message) + if np.isscalar(integer_supercell): + return np.full(3, integer_supercell) + if integer_supercell.shape == (3,): + return integer_supercell + message = ( + f"supercell='{supercell}' is not a scalar or a three component vector." + ) + raise exception.IncorrectUsage(message) + + def _cell_and_transformation(self, standard_form): + if standard_form: + cell_obj = ase.cell.Cell(self.lattice_vectors()) + cell_obj, transformation = cell_obj.standard_form() + cell_string = f"""\ +0.0 {self._format_number(cell_obj[0,0])} xlo xhi +0.0 {self._format_number(cell_obj[1,1])} ylo yhi +0.0 {self._format_number(cell_obj[2,2])} zlo zhi +{self._format_number((cell_obj[1,0], cell_obj[2,0], cell_obj[2,1]))} xy xz yz""" + else: + lattice_vectors = self.lattice_vectors() + cell_string = f"""\ +{self._format_number(lattice_vectors[0])} avec +{self._format_number(lattice_vectors[1])} bvec +{self._format_number(lattice_vectors[2])} cvec +0.0 0.0 0.0 abc origin""" + transformation = np.eye(3) + return cell_string, transformation + + def _position_lines(self, number_ion_types, transformation): + positions = self.cartesian_positions() @ transformation.T + ion_type_labels = [ + str(ion_type + 1) + for ion_type, number in enumerate(number_ion_types) + for _ in range(number) + ] + return "\n".join( + f"{i + 1} {ion_type_labels[i]} {self._format_number(position)}" + for i, position in enumerate(positions) + ) + + def _format_number(self, number): + number = np.atleast_1d(number) + return " ".join(f"{x:24.16E}" for x in number) + + @property + def _is_trajectory(self): + return self._raw_structure.positions.ndim == 3 + + +@dataclass +class _Format: + begin_table: str = "" + column_separator: str = " " + row_separator: str = "\n" + end_table: str = "" + newline: str = "" + + def comment_line(self, stoichiometry, step_string, ion_types): + return f"{stoichiometry.to_string(ion_types)}{step_string}{self.newline}" + + def scaling_factor(self, scale): + return f"{self._element_to_string(scale)}{self.newline}".lstrip() + + def ion_list(self, stoichiometry, ion_types): + return f"{stoichiometry.to_POSCAR(self.newline, ion_types)}{self.newline}" + + def coordinate_system(self): + return f"Direct{self.newline}" + + def vectors_to_table(self, vectors): + rows = (self._vector_to_row(vector) for vector in vectors) + return f"{self.begin_table}{self.row_separator.join(rows)}{self.end_table}" + + def _vector_to_row(self, vector): + elements = (self._element_to_string(element) for element in vector) + return self.column_separator.join(elements) + + def _element_to_string(self, element): + return f"{element:21.16f}" + + +@quantity("structure") +class Structure(view.Mixin): + """The structure contains the unit cell and the position of all ions within. + + The crystal structure is the specific arrangement of ions in a three-dimensional + repeating pattern. This spatial arrangement is characterized by the unit cell and + the relative position of the ions. The unit cell is repeated periodically in three + dimensions to form the crystal. The combination of unit cell and ion positions + determines the symmetry of the crystal. This symmetry helps understanding the + material properties because some symmetries do not allow for the presence of some + properties, e.g., you cannot observe a ferroelectric :data:`~py4vasp.calculation.polarization` + in a system with inversion symmetry. Therefore relaxing the crystal structure with + VASP is an important first step in analyzing materials properties. + + When you run a relaxation or MD simulation, this class allows to access all + individual steps of the trajectory. Typically, you would study the converged + structure after an ionic relaxation or to visualize the changes of the structure + along the simulation. Moreover, you could take snapshots along the trajectory + and further process them by computing more properties. + + Examples + -------- + Let us create some example data so that we can illustrate how to use this class. + Of course you can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you access the structure, the result will depend on the steps that you selected + with the [] operator. Without any selection the results from the final step will be + used. + + >>> calculation.structure.number_steps() + 1 + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.structure[:].number_steps() + 4 + + You can also select specific {step}s or a subset of {step}s as follows + + >>> calculation.structure[3].number_steps() + 1 + >>> calculation.structure[1:4].number_steps() + 3 + """ + + A_to_nm = 0.1 + "Converting Å to nm used for mdtraj trajectories." + + def __init__(self, source, quantity_name="structure", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_structure) -> "Structure": + return cls(source=DataSource(raw_structure)) + + def __repr__(self): + return f"{type(self).__name__}.from_path({self._path!r})" + + def _handler_factory(self, raw): + return StructureHandler.from_data(raw, steps=self._steps) + + def __getitem__(self, steps) -> "Structure": + with self._source.access(self._quantity_name) as raw_structure: + is_trajectory = raw_structure.positions.ndim == 3 + if not is_trajectory: + message = ( + "The structure is not a Trajectory so accessing individual " + "elements is not allowed." + ) + raise exception.IncorrectUsage(message) + new = copy.copy(self) + new._steps = steps + return new + + @classmethod + def from_POSCAR(cls, poscar, *, elements=None): + """Generate a structure from string in POSCAR format. + + The POSCAR format is the standard format to represent crystal structures in + VASP. This method allows to create a structure from a POSCAR string. + To read more about the POSCAR format, please refer to the `VASP manual `_. + + Parameters + ---------- + elements : list[str] + Name of the elements in the order they appear in the POSCAR file. If the + elements are specified in the POSCAR file, this argument is optional and + if set it will overwrite the choice in the POSCAR file. Old POSCAR files + do not specify the name of the elements; in that case this argument is + required. + + Examples + -------- + We can create a GaAs structure from a POSCAR string as follows + + >>> poscar = '''\\ + ... GaAs + ... 5.65325 + ... 0.0 0.5 0.5 + ... 0.5 0.0 0.5 + ... 0.5 0.5 0.0 + ... 1 1 + ... fractional + ... 0.0 0.0 0.0 + ... 0.25 0.25 0.25''' + >>> structure = py4vasp.calculation.structure.from_POSCAR(poscar, elements=['Ga', 'As']) + >>> print(structure.to_POSCAR()) + GaAs + 5.6532... + 0.0... 0.5... 0.5... + 0.5... 0.0... 0.5... + 0.5... 0.5... 0.0... + Ga As + 1 1 + Direct + 0.00... 0.00... 0.00... + 0.25... 0.25... 0.25... + """ + poscar = _replace_or_set_elements(str(poscar), elements) + poscar = parse.POSCAR(poscar) + return cls.from_data(poscar.structure) + + @classmethod + def from_ase(cls, structure): + """Generate a structure from the ase Atoms class.""" + structure = raw.Structure( + stoichiometry=_stoichiometry.raw_stoichiometry_from_ase(structure), + cell=_cell_from_ase(structure), + positions=structure.get_scaled_positions()[np.newaxis], + ) + return cls.from_data(structure) + + def __str__(self, selection=None): + "Generate a string representing the final structure usable as a POSCAR file." + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + StructureHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def _repr_html_(self): + return merge_strings( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler._repr_html_, + ) + + def read(self, ion_types=None): + """Read the structural information into a dictionary. + + The returned dictionary contains the following keys: + - 'lattice_vectors': The lattice vectors of the unit cell. + - 'positions': The positions of the atoms in the unit cell. + - 'elements': The chemical elements of the atoms in the unit cell. + - 'names': The names of the atoms in the unit cell. + + Note that 'elements' and 'names' have the same length as the number of atoms in + the unit cell. + + Parameters + ---------- + ion_types : Sequence + Overwrite the ion types present in the raw data. You can use this to quickly + generate different stoichiometries without modifying the underlying raw data. + + Returns + ------- + dict + Contains the unit cell of the crystal, as well as the position of + all the atoms in units of the lattice vectors and the elements of + the atoms for all selected steps. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `read` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.read() + {'lattice_vectors': array([[...]]), 'positions': array([[...]]), + 'elements': [...], 'names': [...]} + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the lattice vectors and positions contain an additional + dimension for the different steps. + + >>> calculation.structure[:].read() + {'lattice_vectors': array([[[...]]]), 'positions': array([[[...]]]), + 'elements': [...], 'names': [...]} + + You can also select specific steps or a subset of steps as follows + + >>> calculation.structure[1].read() + {'lattice_vectors': array([[...]]), 'positions': array([[...]]), + 'elements': [...], 'names': [...]} + >>> calculation.structure[0:2].read() + {'lattice_vectors': array([[[...]]]), 'positions': array([[[...]]]), + 'elements': [...], 'names': [...]} + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.to_dict, + ion_types, + ) + + def to_dict(self, ion_types=None): + """Convenient alias for :py:meth:`read`. Please read the documentation there.""" + return self.read(ion_types=ion_types) + + def to_view(self, supercell=None, ion_types=None): + """Generate a 3d representation of the structure(s). + + This method uses the `View` class to create a 3d visualization of the atomic + structure(s) in the unit cell. + + Parameters + ---------- + supercell : int or np.ndarray + If present the structure is replicated the specified number of times + along each direction. + ion_types : Sequence + Overwrite the ion types present in the raw data. You can use this to quickly + generate different stoichiometries without modifying the underlying raw data. + + Returns + ------- + View + Visualize the structure(s) as a 3d figure. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_view` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.to_view() + View(elements=array([[...]], dtype=...), lattice_vectors=array([[[...]]]), + positions=array([[[...]]]), ...) + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the lattice vectors and positions contain an additional + dimension for the different steps. + + >>> calculation.structure[:].to_view() + View(elements=array([[...], ..., [...]], dtype=...), lattice_vectors=array([[[...]], ..., [[...]]]), + positions=array([[[...]], ..., [[...]]]), ...) + + You can also select specific steps or a subset of steps as follows + + >>> calculation.structure[1].to_view() + View(elements=array([[...]], dtype=...), lattice_vectors=array([[[...]]]), + positions=array([[[...]]]), ...) + >>> calculation.structure[0:2].to_view() + View(elements=array([[...], [...]], dtype=...), lattice_vectors=array([[[...]], [[...]]]), + positions=array([[[...]], [[...]]]), ...) + + You may also replicate the structure by specifying a supercell. + + >>> calculation.structure.to_view(supercell=2) + View(..., supercell=array([2, 2, 2]), ...) + + The supercell size can also be different for the different directions. + + >>> calculation.structure.to_view(supercell=[2,3,1]) + View(..., supercell=array([2, 3, 1]), ...) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.to_view, + supercell, + ion_types, + ) + + def to_ase(self, supercell=None, ion_types=None): + """Convert the structure to an ASE Atoms object. + + ASE (the Atomic Simulation Environment) is a popular Python package for atomistic + simulations. This method converts the VASP structure to an ASE Atoms object, + which can be used for further analysis and visualization. + + Parameters + ---------- + supercell : int or np.ndarray + If present the structure is replicated the specified number of times + along each direction. + ion_types : Sequence + Overwrite the ion types present in the raw data. You can use this to quickly + generate different stoichiometries without modifying the underlying raw data. + + Returns + ------- + Atoms + Structural information for ASE package. Read more about ASE `here `_. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_ase` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.to_ase() + Atoms(symbols='...', pbc=True, cell=[[...]]) + + You can also select specific steps as follows + + >>> calculation.structure[1].to_ase() + Atoms(symbols='...', pbc=True, cell=[[...]]) + + Notice that converting multiple steps to ASE trajectories is not implemented. + + You may also replicate the structure by specifying a supercell. If you compare + the cell size with the previous example, you will see that it is doubled in all + directions. + + >>> calculation.structure.to_ase(supercell=2) + Atoms(symbols='...', pbc=True, cell=[[...]]) + + The supercell size can also be different for the different directions. The three + lattice vectors will be scaled accordingly. + + >>> calculation.structure.to_ase(supercell=[2,3,1]) + Atoms(symbols='...', pbc=True, cell=[[...]]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.to_ase, + supercell, + ion_types, + ) + + def to_mdtraj(self, ion_types=None): + """Convert the trajectory to mdtraj.Trajectory + + mdtraj is a popular Python package to analyze molecular dynamics trajectories. + This method converts the VASP structure trajectory to an mdtraj.Trajectory + object, which can be used for further analysis and visualization. + + Parameters + ---------- + ion_types : Sequence + Overwrite the ion types present in the raw data. You can use this to quickly + generate different stoichiometries without modifying the underlying raw data. + + Returns + ------- + mdtraj.Trajectory + The mdtraj package offers many functionalities to analyze a MD + trajectory. By converting the VASP data to their format, we facilitate + using all functions of that package. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + To convert the whole trajectory (all steps), you don't specify the array boundaries. + + >>> calculation.structure[:].to_mdtraj() + + + You can also select a subset of steps as follows + + >>> calculation.structure[0:2].to_mdtraj() + + + You cannot convert a single structure to mdtraj.Trajectory. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.to_mdtraj, + ion_types, + ) + + def to_POSCAR(self, ion_types=None): + """Convert the structure(s) to a POSCAR format. + + Use this method to generate a string in POSCAR format representing the + structure(s). You can use this string to write a POSCAR file for VASP. This + can be useful if you want to use the relaxed structure from a VASP calculation + or a snapshot from an MD simulation as input for a new VASP calculation. + + Parameters + ---------- + ion_types : Sequence + Overwrite the ion types present in the raw data. You can use this to quickly + generate different stoichiometries without modifying the underlying raw data. + + Returns + ------- + str + Returns the POSCAR of the selected steps. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_POSCAR` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> poscar = calculation.structure.to_POSCAR() + >>> assert poscar == str(calculation.structure) + + You can also select specific steps as follows + + >>> poscar = calculation.structure[1].to_POSCAR() + >>> assert poscar == str(calculation.structure[1]) + + Notice that converting multiple steps to POSCAR format is not implemented. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.to_POSCAR, + ion_types, + ) + + def to_lammps(self, standard_form=True): + """Convert the structure to LAMMPS format. + + LAMMPS is a popular molecular dynamics simulation software. This method + converts the structure to a string in LAMMPS format, which can be used as + input for LAMMPS simulations. + + Parameters + ---------- + standard_form : bool + Determines whether the structure is standardize, i.e., the lattice vectors + are a triagonal matrix. + + Returns + ------- + str + Returns a string describing the structure for LAMMPS + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_lammps` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> print(calculation.structure.to_lammps()) + Configuration 1: system "..." + ... atoms + ... atom types + ... xlo xhi + ... ylo yhi + ... zlo zhi + ... xy xz yz + Atoms # atomic + 1 1 ... + + You can also select specific steps as follows + + >>> print(calculation.structure[1].to_lammps()) + Configuration 1: system "..." + ... atoms + ... atom types + ... xlo xhi + ... ylo yhi + ... zlo zhi + ... xy xz yz + Atoms # atomic + 1 1 ... + + Notice that converting multiple steps to LAMMPS format is not implemented. + + LAMMPS requires either a standard form of the unit cell or the transformation + from the original cell to the standard form. By default, the standard form is + used. You can disable this behavior as follows + + >>> print(calculation.structure.to_lammps(standard_form=False)) + Configuration 1: system "..." + ... atoms + ... atom types + ... avec + ... bvec + ... cvec + ... abc origin + Atoms # atomic + 1 1 ... + + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.to_lammps, + standard_form, + ) + + def lattice_vectors(self): + """Return the lattice vectors spanning the unit cell + + Returns + ------- + np.ndarray + Lattice vectors of the unit cell in Å. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `lattice_vectors` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.lattice_vectors() + array([[...], [...], [...]]) + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.structure[:].lattice_vectors() + array([[[...]], [[...]], [[...]], [[...]]]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.lattice_vectors, + ) + + def positions(self): + """Return the direct coordinates of all ions in the unit cell. + + Direct or fractional coordinates measure the position of the ions in terms of + the lattice vectors. Hence they are dimensionless quantities. + + Returns + ------- + np.ndarray + Positions of all ions in terms of the lattice vectors. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `positions` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.positions() + array([[...]]) + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.structure[:].positions() + array([[[...]]]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.positions, + ) + + def cartesian_positions(self): + """Convert the positions from direct coordinates to cartesian ones. + + Returns + ------- + np.ndarray + Position of all atoms in cartesian coordinates in Å. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `cartesian_positions` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.cartesian_positions() + array([[...]]) + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.structure[:].cartesian_positions() + array([[[...]]]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.cartesian_positions, + ) + + def volume(self): + """Return the volume of the unit cell for the selected steps. + + Returns + ------- + float or np.ndarray + The volume(s) of the selected step(s) in ų. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this method. + You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `volume` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.structure.volume() + np.float... + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.structure[:].volume() + array([...]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.volume, + ) + + def number_atoms(self): + """Return the total number of atoms in the structure.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.number_atoms, + ) + + def number_steps(self): + """Return the number of structures in the trajectory.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + StructureHandler.number_steps, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + StructureHandler.from_data, + StructureHandler.to_database, + ) + + +def _cell_from_ase(structure): + lattice_vectors = np.array([structure.get_cell()]) + return raw.Cell(lattice_vectors, scale=raw.VaspData(1.0)) + + +def _replace_or_set_elements(poscar, elements): + line_with_elements = 5 + elements = "" if not elements else " ".join(elements) + lines = poscar.split("\n") + if _elements_not_in_poscar(lines[line_with_elements]): + _raise_error_if_elements_not_set(elements) + lines.insert(line_with_elements, elements) + elif elements: + lines[line_with_elements] = elements + return "\n".join(lines) + + +def _elements_not_in_poscar(elements): + elements = elements.split() + return any(element.isdecimal() for element in elements) + + +def _raise_error_if_elements_not_set(elements): + if not elements: + message = """The POSCAR file does not specify the elements needed to create a + Structure. Please pass `elements=[...]` to the `from_POSCAR` routine where + ... are the elements in the same order as in the POSCAR.""" + raise exception.IncorrectUsage(message) + + +class Mixin: + @property + def _structure(self): + return Structure.from_data(self._raw_data.structure) diff --git a/src/py4vasp/_calculation/velocity.py b/src/py4vasp/_calculation/velocity.py index 5862e50b7..a36313160 100644 --- a/src/py4vasp/_calculation/velocity.py +++ b/src/py4vasp/_calculation/velocity.py @@ -1,392 +1,392 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import copy - -import numpy as np - -from py4vasp import _config, raw -from py4vasp._calculation import slice_ -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_strings, - merge_to_database, - quantity, - slice_steps, -) -from py4vasp._calculation.structure import StructureHandler -from py4vasp._raw.data_db import Velocity_DB -from py4vasp._third_party import view - - -class VelocityHandler: - """Handler for velocity data — performs all data access and transformation logic.""" - - velocity_rescale = 200 - - def __init__(self, raw_velocity: raw.Velocity, steps=None): - self._raw_velocity = raw_velocity - self._steps = steps - - @classmethod - def from_data(cls, raw_velocity: raw.Velocity, steps=None) -> "VelocityHandler": - return cls(raw_velocity, steps=steps) - - def __str__(self) -> str: - step = self._last_step - structure = StructureHandler.from_data(self._raw_velocity.structure, steps=step) - velocities = self.to_numpy() - if velocities.ndim == 3: - velocities = velocities[-1] - velocities_str = self._vectors_to_string(velocities) - return f"{structure}\n\n{velocities_str}" - - def _vectors_to_string(self, vectors): - return "\n".join(self._vector_to_string(vector) for vector in vectors) - - def _vector_to_string(self, vector): - return " ".join(self._element_to_string(element) for element in vector) - - def _element_to_string(self, element): - return f"{element:21.16f}" - - def to_dict(self) -> dict: - """Return the structure and ion velocities in a dictionary. - - Returns - ------- - dict - The dictionary contains the ion velocities as well as the structural - information for reference. - """ - structure = StructureHandler.from_data( - self._raw_velocity.structure, steps=self._steps - ) - return { - "structure": structure.to_dict(), - "velocities": self.to_numpy(), - } - - def to_numpy(self) -> np.ndarray: - """Convert the ion velocities for the selected steps into a numpy array.""" - return slice_steps( - np.array(self._raw_velocity.velocities), self._steps, default_ndim=2 - ) - - def to_database(self) -> dict: - """Serialize velocity statistics to the database format.""" - velocities = np.array(self._raw_velocity.velocities) - if velocities.ndim == 2: - final_velocity_norms = np.linalg.norm(velocities, axis=-1) - initial_velocity_norms = final_velocity_norms.copy() - else: - final_velocity_norms = np.linalg.norm(velocities[-1], axis=-1) - initial_velocity_norms = np.linalg.norm(velocities[0], axis=-1) - return Velocity_DB( - final_velocity_min=float(np.min(final_velocity_norms)), - final_velocity_max=float(np.max(final_velocity_norms)), - final_velocity_mean=float(np.mean(final_velocity_norms)), - final_velocity_std=( - float(np.std(final_velocity_norms)) - if len(final_velocity_norms) > 1 - else 0.0 - ), - final_velocity_median=float(np.median(final_velocity_norms)), - final_index_velocity_max=int(np.argmax(final_velocity_norms)), - initial_velocity_min=float(np.min(initial_velocity_norms)), - initial_velocity_max=float(np.max(initial_velocity_norms)), - initial_index_velocity_max=int(np.argmax(initial_velocity_norms)), - ) - - def to_view(self, supercell=None): - """Plot the velocities as vectors in the structure.""" - structure = StructureHandler.from_data(self._raw_velocity.structure) - viewer = structure.to_view(supercell) - velocities = self.velocity_rescale * self.to_numpy() - if velocities.ndim == 2: - velocities = velocities[np.newaxis] - ion_arrow = view.IonArrow( - quantity=velocities, - label="velocities", - color=_config.VASP_COLORS["gray"], - radius=0.2, - ) - viewer.ion_arrows = [ion_arrow] - return viewer - - def number_steps(self) -> int: - """Return the number of velocities in the trajectory.""" - n = len(np.array(self._raw_velocity.velocities)) - return len(range(n)[self._to_slice]) - - @property - def _last_step(self): - if self._steps is None or self._steps == -1: - return -1 - if isinstance(self._steps, slice): - stop = self._steps.stop - return (stop - 1) if stop is not None else -1 - return self._steps - - @property - def _to_slice(self): - if self._steps is None or self._steps == -1: - return slice(-1, None) - if isinstance(self._steps, slice): - return self._steps - return slice(self._steps, self._steps + 1) - - -@quantity("velocity") -class Velocity(view.Mixin): - """The velocities describe the ionic motion during an MD simulation. - - The velocities of the ions are a metric for the temperature of the system. Most - of the time, it is not necessary to consider them explicitly. VASP will set the - velocities automatically according to the temperature settings (:tag:`TEBEG` and - :tag:`TEEND`) unless you set them explicitly in the POSCAR file. Since the - velocities are not something you typically need, VASP will only store them during - the simulation if you set :tag:`VELOCITY` = T in the INCAR file. In that case you - can read the velocities of each step along the trajectory. - - Examples - -------- - Let us create some example data so that we can illustrate how to use this class. - Of course you can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you access the velocities, the result will depend on the steps that you selected - with the [] operator. Without any selection the results from the final step will be - used. - - >>> calculation.velocity.number_steps() - 1 - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.velocity[:].number_steps() - 4 - - You can also select specific steps or a subset of steps as follows - - >>> calculation.velocity[3].number_steps() - 1 - >>> calculation.velocity[1:4].number_steps() - 3 - """ - - velocity_rescale = VelocityHandler.velocity_rescale - - def __init__(self, source, quantity_name: str = "velocity", steps=None): - self._source = source - self._quantity_name = quantity_name - self._steps = steps - - @classmethod - def from_data(cls, raw_velocity: raw.Velocity) -> "Velocity": - """Create a Velocity dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_velocity)) - - def __getitem__(self, steps) -> "Velocity": - new = copy.copy(self) - new._steps = steps - return new - - def _handler_factory(self, raw): - return VelocityHandler.from_data(raw, steps=self._steps) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - VelocityHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self) if not cycle else "...") - - def to_dict(self, selection: str | None = None) -> dict: - """Convenient alias for :py:meth:`read`.""" - return self.read() - - def read(self) -> dict: - """Return the structure and ion velocities in a dictionary. - - Returns - ------- - dict - The dictionary contains the ion velocities as well as the structural - information for reference. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this - method. You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `read` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. The structure is included to provide the necessary context - for the velocities. - - >>> calculation.velocity.read() - {'structure': {...}, 'velocities': array([[...]])} - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the velocities contain an additional dimension for the - different steps. - - >>> calculation.velocity[:].read() - {'structure': {...}, 'velocities': array([[[...]]])} - - You can also select specific steps or a subset of steps as follows - - >>> calculation.velocity[1].read() - {'structure': {...}, 'velocities': array([[...]])} - >>> calculation.velocity[0:2].read() - {'structure': {...}, 'velocities': array([[[...]]])} - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - VelocityHandler.to_dict, - ) - - def to_numpy(self) -> np.ndarray: - """Convert the ion velocities for the selected steps into a numpy array. - - The velocities are given in units of Å/fs. - - Returns - ------- - np.ndarray - A numpy array of the velocities of the selected steps. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this - method. You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_numpy` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.velocity.to_numpy() - array([[...]]) - - To select the results for all steps, you don't specify the array boundaries. - Notice that in this case the velocities contain an additional dimension for the - different steps. - - >>> calculation.velocity[:].to_numpy() - array([[[...]]]) - - You can also select specific steps or a subset of steps as follows - - >>> calculation.velocity[1].to_numpy() - array([[...]]) - >>> calculation.velocity[0:2].to_numpy() - array([[[...], [...]]]) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - VelocityHandler.to_numpy, - ) - - def to_view(self, supercell=None) -> view.View: - """Plot the velocities as vectors in the structure. - - This method adds arrows to the atoms in the structure sized according to the - size of the velocity. The length of the arrows is scaled by a constant - `velocity_rescale` to convert from Å/fs to a length in Å. - - Parameters - ---------- - supercell : int or np.ndarray - If present the structure is replicated the specified number of times - along each direction. - - Returns - ------- - View - Contains all atoms and the velocities are drawn as vectors. - - Examples - -------- - First, we create some example data so that we can illustrate how to use this - method. You can also use your own VASP calculation data if you have it available. - - >>> from py4vasp import demo - >>> calculation = demo.calculation(path) - - If you use the `to_view` method, the result will depend on the steps that you - selected with the [] operator. Without any selection the results from the final - step will be used. - - >>> calculation.velocity.to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='velocities', ...)], ...) - - To select the results for all steps, you don't specify the array boundaries. - - >>> calculation.velocity[:].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='velocities', ...)], ...) - - You can also select specific steps or a subset of steps as follows - - >>> calculation.velocity[1].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='velocities', ...)], ...) - >>> calculation.velocity[0:2].to_view() - View(..., ion_arrows=[IonArrow(quantity=array([[[...], [...]]]), label='velocities', ...)], ...) - - You may also replicate the structure by specifying a supercell. - - >>> calculation.velocity.to_view(supercell=2) - View(..., supercell=array([2, 2, 2]), ...) - - The supercell size can also be different for the different directions. - - >>> calculation.velocity.to_view(supercell=[2,3,1]) - View(..., supercell=array([2, 3, 1]), ...) - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - VelocityHandler.to_view, - supercell, - ) - - def number_steps(self) -> int: - """Return the number of velocities in the trajectory.""" - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - VelocityHandler.number_steps, - ) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - VelocityHandler.from_data, - VelocityHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import copy + +import numpy as np + +from py4vasp import _config, raw +from py4vasp._calculation import slice_ +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_strings, + merge_to_database, + quantity, + slice_steps, +) +from py4vasp._calculation.structure import StructureHandler +from py4vasp._raw.data_db import Velocity_DB +from py4vasp._third_party import view + + +class VelocityHandler: + """Handler for velocity data — performs all data access and transformation logic.""" + + velocity_rescale = 200 + + def __init__(self, raw_velocity: raw.Velocity, steps=None): + self._raw_velocity = raw_velocity + self._steps = steps + + @classmethod + def from_data(cls, raw_velocity: raw.Velocity, steps=None) -> "VelocityHandler": + return cls(raw_velocity, steps=steps) + + def __str__(self) -> str: + step = self._last_step + structure = StructureHandler.from_data(self._raw_velocity.structure, steps=step) + velocities = self.to_numpy() + if velocities.ndim == 3: + velocities = velocities[-1] + velocities_str = self._vectors_to_string(velocities) + return f"{structure}\n\n{velocities_str}" + + def _vectors_to_string(self, vectors): + return "\n".join(self._vector_to_string(vector) for vector in vectors) + + def _vector_to_string(self, vector): + return " ".join(self._element_to_string(element) for element in vector) + + def _element_to_string(self, element): + return f"{element:21.16f}" + + def to_dict(self) -> dict: + """Return the structure and ion velocities in a dictionary. + + Returns + ------- + dict + The dictionary contains the ion velocities as well as the structural + information for reference. + """ + structure = StructureHandler.from_data( + self._raw_velocity.structure, steps=self._steps + ) + return { + "structure": structure.to_dict(), + "velocities": self.to_numpy(), + } + + def to_numpy(self) -> np.ndarray: + """Convert the ion velocities for the selected steps into a numpy array.""" + return slice_steps( + np.array(self._raw_velocity.velocities), self._steps, default_ndim=2 + ) + + def to_database(self) -> dict: + """Serialize velocity statistics to the database format.""" + velocities = np.array(self._raw_velocity.velocities) + if velocities.ndim == 2: + final_velocity_norms = np.linalg.norm(velocities, axis=-1) + initial_velocity_norms = final_velocity_norms.copy() + else: + final_velocity_norms = np.linalg.norm(velocities[-1], axis=-1) + initial_velocity_norms = np.linalg.norm(velocities[0], axis=-1) + return Velocity_DB( + final_velocity_min=float(np.min(final_velocity_norms)), + final_velocity_max=float(np.max(final_velocity_norms)), + final_velocity_mean=float(np.mean(final_velocity_norms)), + final_velocity_std=( + float(np.std(final_velocity_norms)) + if len(final_velocity_norms) > 1 + else 0.0 + ), + final_velocity_median=float(np.median(final_velocity_norms)), + final_index_velocity_max=int(np.argmax(final_velocity_norms)), + initial_velocity_min=float(np.min(initial_velocity_norms)), + initial_velocity_max=float(np.max(initial_velocity_norms)), + initial_index_velocity_max=int(np.argmax(initial_velocity_norms)), + ) + + def to_view(self, supercell=None): + """Plot the velocities as vectors in the structure.""" + structure = StructureHandler.from_data(self._raw_velocity.structure) + viewer = structure.to_view(supercell) + velocities = self.velocity_rescale * self.to_numpy() + if velocities.ndim == 2: + velocities = velocities[np.newaxis] + ion_arrow = view.IonArrow( + quantity=velocities, + label="velocities", + color=_config.VASP_COLORS["gray"], + radius=0.2, + ) + viewer.ion_arrows = [ion_arrow] + return viewer + + def number_steps(self) -> int: + """Return the number of velocities in the trajectory.""" + n = len(np.array(self._raw_velocity.velocities)) + return len(range(n)[self._to_slice]) + + @property + def _last_step(self): + if self._steps is None or self._steps == -1: + return -1 + if isinstance(self._steps, slice): + stop = self._steps.stop + return (stop - 1) if stop is not None else -1 + return self._steps + + @property + def _to_slice(self): + if self._steps is None or self._steps == -1: + return slice(-1, None) + if isinstance(self._steps, slice): + return self._steps + return slice(self._steps, self._steps + 1) + + +@quantity("velocity") +class Velocity(view.Mixin): + """The velocities describe the ionic motion during an MD simulation. + + The velocities of the ions are a metric for the temperature of the system. Most + of the time, it is not necessary to consider them explicitly. VASP will set the + velocities automatically according to the temperature settings (:tag:`TEBEG` and + :tag:`TEEND`) unless you set them explicitly in the POSCAR file. Since the + velocities are not something you typically need, VASP will only store them during + the simulation if you set :tag:`VELOCITY` = T in the INCAR file. In that case you + can read the velocities of each step along the trajectory. + + Examples + -------- + Let us create some example data so that we can illustrate how to use this class. + Of course you can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you access the velocities, the result will depend on the steps that you selected + with the [] operator. Without any selection the results from the final step will be + used. + + >>> calculation.velocity.number_steps() + 1 + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.velocity[:].number_steps() + 4 + + You can also select specific steps or a subset of steps as follows + + >>> calculation.velocity[3].number_steps() + 1 + >>> calculation.velocity[1:4].number_steps() + 3 + """ + + velocity_rescale = VelocityHandler.velocity_rescale + + def __init__(self, source, quantity_name: str = "velocity", steps=None): + self._source = source + self._quantity_name = quantity_name + self._steps = steps + + @classmethod + def from_data(cls, raw_velocity: raw.Velocity) -> "Velocity": + """Create a Velocity dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_velocity)) + + def __getitem__(self, steps) -> "Velocity": + new = copy.copy(self) + new._steps = steps + return new + + def _handler_factory(self, raw): + return VelocityHandler.from_data(raw, steps=self._steps) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + VelocityHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self) if not cycle else "...") + + def to_dict(self, selection: str | None = None) -> dict: + """Convenient alias for :py:meth:`read`.""" + return self.read() + + def read(self) -> dict: + """Return the structure and ion velocities in a dictionary. + + Returns + ------- + dict + The dictionary contains the ion velocities as well as the structural + information for reference. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this + method. You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `read` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. The structure is included to provide the necessary context + for the velocities. + + >>> calculation.velocity.read() + {'structure': {...}, 'velocities': array([[...]])} + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the velocities contain an additional dimension for the + different steps. + + >>> calculation.velocity[:].read() + {'structure': {...}, 'velocities': array([[[...]]])} + + You can also select specific steps or a subset of steps as follows + + >>> calculation.velocity[1].read() + {'structure': {...}, 'velocities': array([[...]])} + >>> calculation.velocity[0:2].read() + {'structure': {...}, 'velocities': array([[[...]]])} + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + VelocityHandler.to_dict, + ) + + def to_numpy(self) -> np.ndarray: + """Convert the ion velocities for the selected steps into a numpy array. + + The velocities are given in units of Å/fs. + + Returns + ------- + np.ndarray + A numpy array of the velocities of the selected steps. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this + method. You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_numpy` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.velocity.to_numpy() + array([[...]]) + + To select the results for all steps, you don't specify the array boundaries. + Notice that in this case the velocities contain an additional dimension for the + different steps. + + >>> calculation.velocity[:].to_numpy() + array([[[...]]]) + + You can also select specific steps or a subset of steps as follows + + >>> calculation.velocity[1].to_numpy() + array([[...]]) + >>> calculation.velocity[0:2].to_numpy() + array([[[...], [...]]]) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + VelocityHandler.to_numpy, + ) + + def to_view(self, supercell=None) -> view.View: + """Plot the velocities as vectors in the structure. + + This method adds arrows to the atoms in the structure sized according to the + size of the velocity. The length of the arrows is scaled by a constant + `velocity_rescale` to convert from Å/fs to a length in Å. + + Parameters + ---------- + supercell : int or np.ndarray + If present the structure is replicated the specified number of times + along each direction. + + Returns + ------- + View + Contains all atoms and the velocities are drawn as vectors. + + Examples + -------- + First, we create some example data so that we can illustrate how to use this + method. You can also use your own VASP calculation data if you have it available. + + >>> from py4vasp import demo + >>> calculation = demo.calculation(path) + + If you use the `to_view` method, the result will depend on the steps that you + selected with the [] operator. Without any selection the results from the final + step will be used. + + >>> calculation.velocity.to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='velocities', ...)], ...) + + To select the results for all steps, you don't specify the array boundaries. + + >>> calculation.velocity[:].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...]]]), label='velocities', ...)], ...) + + You can also select specific steps or a subset of steps as follows + + >>> calculation.velocity[1].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[...]]), label='velocities', ...)], ...) + >>> calculation.velocity[0:2].to_view() + View(..., ion_arrows=[IonArrow(quantity=array([[[...], [...]]]), label='velocities', ...)], ...) + + You may also replicate the structure by specifying a supercell. + + >>> calculation.velocity.to_view(supercell=2) + View(..., supercell=array([2, 2, 2]), ...) + + The supercell size can also be different for the different directions. + + >>> calculation.velocity.to_view(supercell=[2,3,1]) + View(..., supercell=array([2, 3, 1]), ...) + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + VelocityHandler.to_view, + supercell, + ) + + def number_steps(self) -> int: + """Return the number of velocities in the trajectory.""" + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + VelocityHandler.number_steps, + ) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + VelocityHandler.from_data, + VelocityHandler.to_database, + ) diff --git a/src/py4vasp/_calculation/workfunction.py b/src/py4vasp/_calculation/workfunction.py index 689e7c1e8..3d4146f3f 100644 --- a/src/py4vasp/_calculation/workfunction.py +++ b/src/py4vasp/_calculation/workfunction.py @@ -1,194 +1,194 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from contextlib import suppress - -from py4vasp import exception, raw -from py4vasp._calculation import bandgap as bandgap_module -from py4vasp._calculation.dispatch import ( - DataSource, - _dispatch, - merge_default, - merge_graphs, - merge_strings, - merge_to_database, - quantity, -) -from py4vasp._raw.data_db import Workfunction_DB -from py4vasp._third_party import graph - - -class WorkfunctionHandler: - """Handler for the workfunction quantity. Works with exactly one raw.Workfunction object.""" - - def __init__(self, raw_workfunction: raw.Workfunction): - self._raw_workfunction = raw_workfunction - - @classmethod - def from_data(cls, raw_workfunction: raw.Workfunction) -> "WorkfunctionHandler": - return cls(raw_workfunction) - - def read(self) -> dict: - """Reports useful information about the workfunction as a dictionary. - - In addition to the vacuum potential, the dictionary contains typical reference - energies such as the valence band maximum, the conduction band minimum, and the - Fermi energy. Furthermore you obtain the average potential, so you can use a - different algorithm to determine the vacuum potential if desired. - - Returns - ------- - dict - Contains vacuum potential, average potential and relevant reference energies - within the surface. - """ - band_extrema = {} - with suppress(exception.NoData): - gap = bandgap_module.BandgapHandler.from_data( - self._raw_workfunction.reference_potential - ) - band_extrema = { - "valence_band_maximum": gap.valence_band_maximum(), - "conduction_band_minimum": gap.conduction_band_minimum(), - } - return { - "direction": f"lattice vector {self._raw_workfunction.idipol}", - "distance": self._raw_workfunction.distance[:], - "average_potential": self._raw_workfunction.average_potential[:], - "vacuum_potential": self._raw_workfunction.vacuum_potential[:], - "fermi_energy": self._raw_workfunction.fermi_energy, - **band_extrema, - } - - def to_dict(self) -> dict: - """Public alias for read().""" - return self.read() - - def to_database(self) -> dict: - """Serialize workfunction data for database storage.""" - return Workfunction_DB( - direction=self._raw_workfunction.idipol, - workfunction_value=None, - ) - - def to_graph(self) -> graph.Graph: - """Plot the average potential along the lattice vector selected by IDIPOL. - - Returns - ------- - Graph - A plot where the distance in the unit cell along the selected lattice vector - is on the x axis and the averaged potential across the plane of the other - two lattice vectors is on the y axis. - """ - data = self.read() - series = graph.Series(data["distance"], data["average_potential"], "potential") - return graph.Graph( - series=series, - xlabel=f"distance along {data['direction']} (Å)", - ylabel="average potential (eV)", - ) - - def __str__(self) -> str: - data = self.read() - return f"""workfunction along {data["direction"]}: - vacuum potential: {data["vacuum_potential"][0]:.3f} {data["vacuum_potential"][1]:.3f} - Fermi energy: {data["fermi_energy"]:.3f} - valence band maximum: {data["valence_band_maximum"]:.3f} - conduction band minimum: {data["conduction_band_minimum"]:.3f}""" - - -@quantity("workfunction") -class Workfunction(graph.Mixin): - """The workfunction describes the energy required to remove an electron to the vacuum. - - The workfunction of a material is the minimum energy required to remove an - electron from its most loosely bound state and move it to an energy level just - outside the material's surface. In other words, it represents the energy barrier - that electrons must overcome to escape the material. The workfunction helps - understanding electronic emission phenomena in surface science and materials - engineering. In VASP, you can compute the workfunction by setting the :tag:`IDIPOL` - flag in the INCAR file. This class provides then the functionality to analyze the - resulting potential. - """ - - def __init__(self, source, quantity_name: str = "workfunction"): - self._source = source - self._quantity_name = quantity_name - - @classmethod - def from_data(cls, raw_workfunction: raw.Workfunction) -> "Workfunction": - """Create a Workfunction dispatcher from raw data (convenience for testing).""" - return cls(source=DataSource(raw_workfunction)) - - @property - def path(self): - """Returns the path from which the output is obtained.""" - return self._path - - def _handler_factory(self, raw_data): - return WorkfunctionHandler.from_data(raw_data) - - def read(self) -> dict: - """Reports useful information about the workfunction as a dictionary. - - In addition to the vacuum potential, the dictionary contains typical reference - energies such as the valence band maximum, the conduction band minimum, and the - Fermi energy. Furthermore you obtain the average potential, so you can use a - different algorithm to determine the vacuum potential if desired. - - Returns - ------- - dict - Contains vacuum potential, average potential and relevant reference energies - within the surface. - """ - return merge_default( - self._source, - self._quantity_name, - None, - self._handler_factory, - WorkfunctionHandler.read, - ) - - def to_dict(self, selection: str | None = None) -> dict: - """Public alias for read(). Check that method for examples and optional arguments.""" - return self.read() - - def to_graph(self) -> graph.Graph: - """Plot the average potential along the lattice vector selected by IDIPOL. - - Returns - ------- - Graph - A plot where the distance in the unit cell along the selected lattice vector - is on the x axis and the averaged potential across the plane of the other - two lattice vectors is on the y axis. - """ - return merge_graphs( - self._source, - self._quantity_name, - None, - self._handler_factory, - WorkfunctionHandler.to_graph, - ) - - def __str__(self, selection=None) -> str: - return merge_strings( - self._source, - self._quantity_name, - selection, - self._handler_factory, - WorkfunctionHandler.__str__, - ) - - def _repr_pretty_(self, p, cycle): - p.text(str(self)) - - def _to_database(self) -> dict: - """Return {quantity[_selection]: handler_result} for database storage.""" - return merge_to_database( - self._source, - self._quantity_name, - WorkfunctionHandler.from_data, - WorkfunctionHandler.to_database, - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from contextlib import suppress + +from py4vasp import exception, raw +from py4vasp._calculation import bandgap as bandgap_module +from py4vasp._calculation.dispatch import ( + DataSource, + _dispatch, + merge_default, + merge_graphs, + merge_strings, + merge_to_database, + quantity, +) +from py4vasp._raw.data_db import Workfunction_DB +from py4vasp._third_party import graph + + +class WorkfunctionHandler: + """Handler for the workfunction quantity. Works with exactly one raw.Workfunction object.""" + + def __init__(self, raw_workfunction: raw.Workfunction): + self._raw_workfunction = raw_workfunction + + @classmethod + def from_data(cls, raw_workfunction: raw.Workfunction) -> "WorkfunctionHandler": + return cls(raw_workfunction) + + def read(self) -> dict: + """Reports useful information about the workfunction as a dictionary. + + In addition to the vacuum potential, the dictionary contains typical reference + energies such as the valence band maximum, the conduction band minimum, and the + Fermi energy. Furthermore you obtain the average potential, so you can use a + different algorithm to determine the vacuum potential if desired. + + Returns + ------- + dict + Contains vacuum potential, average potential and relevant reference energies + within the surface. + """ + band_extrema = {} + with suppress(exception.NoData): + gap = bandgap_module.BandgapHandler.from_data( + self._raw_workfunction.reference_potential + ) + band_extrema = { + "valence_band_maximum": gap.valence_band_maximum(), + "conduction_band_minimum": gap.conduction_band_minimum(), + } + return { + "direction": f"lattice vector {self._raw_workfunction.idipol}", + "distance": self._raw_workfunction.distance[:], + "average_potential": self._raw_workfunction.average_potential[:], + "vacuum_potential": self._raw_workfunction.vacuum_potential[:], + "fermi_energy": self._raw_workfunction.fermi_energy, + **band_extrema, + } + + def to_dict(self) -> dict: + """Public alias for read().""" + return self.read() + + def to_database(self) -> dict: + """Serialize workfunction data for database storage.""" + return Workfunction_DB( + direction=self._raw_workfunction.idipol, + workfunction_value=None, + ) + + def to_graph(self) -> graph.Graph: + """Plot the average potential along the lattice vector selected by IDIPOL. + + Returns + ------- + Graph + A plot where the distance in the unit cell along the selected lattice vector + is on the x axis and the averaged potential across the plane of the other + two lattice vectors is on the y axis. + """ + data = self.read() + series = graph.Series(data["distance"], data["average_potential"], "potential") + return graph.Graph( + series=series, + xlabel=f"distance along {data['direction']} (Å)", + ylabel="average potential (eV)", + ) + + def __str__(self) -> str: + data = self.read() + return f"""workfunction along {data["direction"]}: + vacuum potential: {data["vacuum_potential"][0]:.3f} {data["vacuum_potential"][1]:.3f} + Fermi energy: {data["fermi_energy"]:.3f} + valence band maximum: {data["valence_band_maximum"]:.3f} + conduction band minimum: {data["conduction_band_minimum"]:.3f}""" + + +@quantity("workfunction") +class Workfunction(graph.Mixin): + """The workfunction describes the energy required to remove an electron to the vacuum. + + The workfunction of a material is the minimum energy required to remove an + electron from its most loosely bound state and move it to an energy level just + outside the material's surface. In other words, it represents the energy barrier + that electrons must overcome to escape the material. The workfunction helps + understanding electronic emission phenomena in surface science and materials + engineering. In VASP, you can compute the workfunction by setting the :tag:`IDIPOL` + flag in the INCAR file. This class provides then the functionality to analyze the + resulting potential. + """ + + def __init__(self, source, quantity_name: str = "workfunction"): + self._source = source + self._quantity_name = quantity_name + + @classmethod + def from_data(cls, raw_workfunction: raw.Workfunction) -> "Workfunction": + """Create a Workfunction dispatcher from raw data (convenience for testing).""" + return cls(source=DataSource(raw_workfunction)) + + @property + def path(self): + """Returns the path from which the output is obtained.""" + return self._path + + def _handler_factory(self, raw_data): + return WorkfunctionHandler.from_data(raw_data) + + def read(self) -> dict: + """Reports useful information about the workfunction as a dictionary. + + In addition to the vacuum potential, the dictionary contains typical reference + energies such as the valence band maximum, the conduction band minimum, and the + Fermi energy. Furthermore you obtain the average potential, so you can use a + different algorithm to determine the vacuum potential if desired. + + Returns + ------- + dict + Contains vacuum potential, average potential and relevant reference energies + within the surface. + """ + return merge_default( + self._source, + self._quantity_name, + None, + self._handler_factory, + WorkfunctionHandler.read, + ) + + def to_dict(self, selection: str | None = None) -> dict: + """Public alias for read(). Check that method for examples and optional arguments.""" + return self.read() + + def to_graph(self) -> graph.Graph: + """Plot the average potential along the lattice vector selected by IDIPOL. + + Returns + ------- + Graph + A plot where the distance in the unit cell along the selected lattice vector + is on the x axis and the averaged potential across the plane of the other + two lattice vectors is on the y axis. + """ + return merge_graphs( + self._source, + self._quantity_name, + None, + self._handler_factory, + WorkfunctionHandler.to_graph, + ) + + def __str__(self, selection=None) -> str: + return merge_strings( + self._source, + self._quantity_name, + selection, + self._handler_factory, + WorkfunctionHandler.__str__, + ) + + def _repr_pretty_(self, p, cycle): + p.text(str(self)) + + def _to_database(self) -> dict: + """Return {quantity[_selection]: handler_result} for database storage.""" + return merge_to_database( + self._source, + self._quantity_name, + WorkfunctionHandler.from_data, + WorkfunctionHandler.to_database, + ) diff --git a/src/py4vasp/_third_party/graph/graph.py b/src/py4vasp/_third_party/graph/graph.py index e6f6a1734..d8862a918 100644 --- a/src/py4vasp/_third_party/graph/graph.py +++ b/src/py4vasp/_third_party/graph/graph.py @@ -15,7 +15,7 @@ from py4vasp._third_party.graph.contour import Contour from py4vasp._third_party.graph.series import Series from py4vasp._third_party.graph.trace import Trace -from py4vasp._util import import_ +from py4vasp._util import import_, merge go = import_.optional("plotly.graph_objects") subplots = import_.optional("plotly.subplots") @@ -564,13 +564,4 @@ def _merge_fields(left_graph, right_graph): def _merge_field(left_graph, right_graph, field_name): left_field = getattr(left_graph, field_name) right_field = getattr(right_graph, field_name) - if not left_field: - return right_field - if not right_field: - return left_field - if left_field != right_field: - message = f"""Cannot combine two graphs with incompatible {field_name}: - left: {left_field} - right: {right_field}""" - raise exception.IncorrectUsage(message) - return left_field + return merge.merge_field_or_raise(left_field, right_field, field_name, "graphs") diff --git a/src/py4vasp/_third_party/view/view.py b/src/py4vasp/_third_party/view/view.py index 0b2988efb..469dc6209 100644 --- a/src/py4vasp/_third_party/view/view.py +++ b/src/py4vasp/_third_party/view/view.py @@ -4,19 +4,19 @@ import os import tempfile from contextlib import suppress -from dataclasses import dataclass +from dataclasses import dataclass, fields from typing import NamedTuple, Optional, Sequence import numpy as np import numpy.typing as npt from py4vasp import exception -from py4vasp._util import convert, import_ +from py4vasp._util import convert, import_, merge ase = import_.optional("ase") ase_cube = import_.optional("ase.io.cube") nglview = import_.optional("nglview") -vaspview = import_.optional("vasp_viewer") +vaspview = import_.optional("vasp.viewer") CUBE_FILENAME = "quantity.cube" @@ -274,6 +274,17 @@ class View: def __post_init__(self): self._verify() + def __add__(self, other): + if not isinstance(other, View): + return NotImplemented + merged = _merge_view_fields(self, other) + # Construct without calling __post_init__: validation is intentionally + # deferred to viewer conversion (to_ngl / to_vasp_viewer). + combined = object.__new__(View) + for field in fields(View): + setattr(combined, field.name, merged[field.name]) + return combined + def _ipython_display_(self, mode: str = "auto"): if mode == "auto": if import_.is_imported(vaspview): @@ -289,7 +300,10 @@ def _ipython_display_(self, mode: str = "auto"): widget = self.to_ngl() widget._ipython_display_() elif mode == "vasp_viewer": + import IPython.display + widget = self.to_vasp_viewer() + IPython.display.display(widget) else: raise exception.IncorrectUsage( f"Mode '{mode}' is not supported. Choose either 'auto', 'ngl' or 'vasp_viewer'." @@ -329,7 +343,18 @@ def to_vasp_viewer(self): This method creates the widget required to view a structure, isosurfaces and arrows at atom centers. The attributes of View are added to a dictionary with which - to call initialize a VASP Viewer widget.""" + to initialize a VASP Viewer widget.""" + structure = self.to_vasp_viewer_config() + widget = vaspview.Widget(structure) + return widget + + def to_vasp_viewer_config(self): + """Create a dictionary with the configuration for VASP Viewer + + This method creates a dictionary with the configuration required to view a + structure, isosurfaces and arrows at atom centers. The attributes of View are + added to a dictionary with which to initialize a VASP Viewer widget. + """ self._verify() structure: dict = { "atoms_trajectory": self._convert_to_list(self.positions), @@ -338,8 +363,7 @@ def to_vasp_viewer(self): } # === Atoms options === - if self.atom_radius is not None: - structure["selections_atom_radius"] = self.atom_radius + structure["selections_atom_radius"] = self.atom_radius or 1.0 # === Vector Group options === if self.ion_arrows is not None: @@ -353,30 +377,44 @@ def to_vasp_viewer(self): for arrow in self.ion_arrows ] if self.grid_scalars is not None: - # TODO merge isosurface branch - # TODO handle list of grid scalars instead of single grid scalar only - # TODO adjust UI to support this - structure["grid_scalar_groups"] = [ - { - "label": grid_quantity.label, - "data": grid_quantity.quantity, # TODO check type - "isosurfaces": [ # TODO hook this list to isosurface settings + # TODO allow time-dependent isosurfaces (viewer-side requirement) + # TODO allow multiple isolevels for the same volume dataset (viewer-side requirement) + structure["volume_datasets"] = [] + for grid_quantity in self.grid_scalars: + if len(grid_quantity.isosurfaces) > 0: + structure["volume_datasets"].extend( { - "isolevel": isosurface.isolevel, - "color": isosurface.color, # TODO interpret this as base color of isosurface - "opacity": isosurface.opacity, # TODO tie this to opacity on isosurface + "label": grid_quantity.label + f" ({idi})", + "data": self._convert_to_list(grid_quantity.quantity[0]), + "grid": grid_quantity.quantity.shape[1:], + "initial_iso_value": isosurface.isolevel, + "color_surface": isosurface.color, } - for isosurface in grid_quantity.isosurfaces - ], - } - for grid_quantity in self.grid_scalars - ] + for idi, isosurface in enumerate(grid_quantity.isosurfaces) + ) + else: + structure["volume_datasets"].append( + { + "label": grid_quantity.label, + "data": self._convert_to_list(grid_quantity.quantity[0]), + "grid": grid_quantity.quantity.shape[1:], + } + ) # === Lattice options === if self.shift is not None: structure["selections_constant_shift"] = self._convert_to_list(self.shift) if self.supercell is not None: - structure["selections_supercell"] = self._convert_to_list(self.supercell) + supercell_list = self._convert_to_list(self.supercell) + if len(supercell_list) == 1: + supercell_list = supercell_list * 3 + elif len(supercell_list) == 2: + if supercell_list[0] != supercell_list[1]: + supercell_list = supercell_list + [1] + else: + supercell_list = supercell_list + [supercell_list[0]] + supercell_list = [0, 0, 0] + supercell_list + structure["selections_bounds"] = supercell_list # === Visualization options === if self.camera is not None: @@ -384,10 +422,7 @@ def to_vasp_viewer(self): if self.show_cell is not None: structure["selections_show_lattice"] = self.show_cell if self.show_axes is not None: - structure["selections_show_xyz"] = False structure["selections_show_abc"] = self.show_axes - structure["selections_show_xyz_aside"] = self.show_axes - structure["selections_show_abc_aside"] = self.show_axes if self.show_axes_at is not None: structure["selections_axes_abc_shift"] = self._convert_to_list( self.show_axes_at @@ -400,7 +435,7 @@ def to_vasp_viewer(self): if self.structure_title: structure["selections_descriptor"] = self.structure_title - return vaspview.Widget(structure) + return structure def _verify(self, mode=None): self._raise_error_if_present_on_multiple_steps(self.grid_scalars, mode) @@ -540,3 +575,50 @@ def _show_arrows_at_atoms(self, widget, trajectory): transformation, ) widget.shape.add_arrow(*(arrow_3d.to_serializable())) + + +def _merge_view_fields(left_view, right_view): + merged = {} + for field in fields(View): + if field.name in ("grid_scalars", "ion_arrows"): + merged[field.name] = _merge_special_sequence( + getattr(left_view, field.name), + getattr(right_view, field.name), + ) + continue + merged[field.name] = _merge_view_field( + left_view, + right_view, + field.name, + ) + return merged + + +def _merge_special_sequence(left_values, right_values): + return merge.merge_unique_sequences(left_values, right_values, _entries_equal) + + +def _entries_equal(left_entry, right_entry): + if left_entry is right_entry: + return True + if type(left_entry) is not type(right_entry): + return False + if hasattr(left_entry, "__dataclass_fields__"): + for field in fields(left_entry): + if not _values_equal( + getattr(left_entry, field.name), + getattr(right_entry, field.name), + ): + return False + return True + return _values_equal(left_entry, right_entry) + + +def _merge_view_field(left_view, right_view, field_name): + left_field = getattr(left_view, field_name) + right_field = getattr(right_view, field_name) + return merge.merge_field_or_raise(left_field, right_field, field_name, "views") + + +def _values_equal(left_value, right_value): + return merge.values_equal(left_value, right_value) diff --git a/src/py4vasp/_util/loadable.py b/src/py4vasp/_util/loadable.py new file mode 100644 index 000000000..5b67bd018 --- /dev/null +++ b/src/py4vasp/_util/loadable.py @@ -0,0 +1,392 @@ +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +"""Utilities to determine loadable quantity selections for a Calculation.""" + +import dataclasses +import inspect + +import h5py +import numpy as np + +from py4vasp import exception, raw +from py4vasp._raw.definition import DEFAULT_FILE, schema, unique_selections +from py4vasp._raw.schema import DEFAULT_SELECTION, Length, Link +from py4vasp._util import check, convert + + +def loadable_sources( + calculation, + call_name, + schema_name, + method, + open_files, + stack, + cache, + legacy_quantities, + errors, +): + """Return loadable selections for one quantity as a list of source names. + + Selections whose data is theoretically available but that fail to load because + of an out-of-memory error are omitted from the returned list. Instead, a + descriptive message is appended to *errors* so the caller can report it without + aborting the inspection of the remaining quantities. + """ + method_name = method or "read" + try: + sources = list(unique_selections(schema_name)) + except exception.FileAccessError: + return [] + if not _implements(calculation, call_name, method_name): + return [] + convention = _source_convention(calculation, call_name, legacy_quantities) + loadable = [] + for source_name in sources: + conv = "plain" if source_name == DEFAULT_SELECTION else convention + if conv is None: + # The source cannot be addressed yet (e.g. a not-yet-migrated quantity). + continue + try: + confirmed = _confirm_selection( + calculation, + call_name, + schema_name, + source_name, + method_name, + conv, + open_files, + stack, + cache, + legacy_quantities, + ) + except MemoryError as error: + errors.append( + _out_of_memory_message(call_name, method_name, source_name, conv, error) + ) + continue + if not confirmed: + continue + loadable.append(source_name) + return loadable + + +def possible_sources(schema_name): + """Return all schema-defined selections for one quantity.""" + try: + return list(unique_selections(schema_name)) + except exception.FileAccessError: + return [] + + +def implements_method(calculation, call_name, method_name): + """Return whether the quantity provides the requested method.""" + return _implements(calculation, call_name, method_name) + + +def _implements(calculation, call_name, method_name): + try: + quantity = _quantity_object(calculation, call_name) + except Exception: + return False + return callable(getattr(quantity, method_name, None)) + + +def _source_convention(calculation, call_name, legacy_quantities): + """Return how a non-default source is addressed: 'keyword', 'index', or None.""" + try: + quantity = _quantity_object(calculation, call_name) + except Exception: + return None + read = getattr(quantity, "read", None) + if read is not None and _accepts_selection(read): + return "keyword" + if call_name not in legacy_quantities and hasattr(type(quantity), "__getitem__"): + return "index" + return None + + +def _confirm_selection( + calculation, + call_name, + schema_name, + source_name, + method_name, + convention, + open_files, + stack, + cache, + legacy_quantities, +): + if method_name == "read": + return _confirm_read( + calculation, + call_name, + schema_name, + source_name, + open_files, + stack, + cache, + legacy_quantities, + ) + # For explicit methods, successful invocation proves the source can be loaded. + if ( + _schema_satisfied( + calculation, + schema_name, + source_name, + open_files, + stack, + cache, + legacy_quantities, + ) + is False + ): + return False + return _invoke( + calculation, + call_name, + method_name, + source_name, + legacy_quantities, + convention, + ) + + +def _confirm_read( + calculation, + call_name, + schema_name, + source_name, + open_files, + stack, + cache, + legacy_quantities, +): + key = (schema_name, source_name) + if key in cache: + return cache[key] + cache[key] = True # provisional value to break cyclic Link dependencies + try: + verdict = _schema_satisfied( + calculation, + schema_name, + source_name, + open_files, + stack, + cache, + legacy_quantities, + ) + if verdict is None: + verdict = _invoke( + calculation, + call_name, + "read", + source_name, + legacy_quantities, + ) + except MemoryError: + # The data is available but too large to load; drop the provisional cache + # entry and let the caller decide how to report the failure. + cache.pop(key, None) + raise + cache[key] = bool(verdict) + return cache[key] + + +def _schema_satisfied( + calculation, + schema_name, + source_name, + open_files, + stack, + cache, + legacy_quantities, +): + try: + source = schema.sources[schema_name][source_name] + except KeyError: + return False + if source.required is not None: + version = _file_version(calculation, open_files, stack) + if version is None or version < source.required: + return False + filename = calculation._file or source.file or DEFAULT_FILE + if source.data is None: + return (calculation._path / filename).exists() + h5f = _open_h5(calculation, open_files, stack, filename) + if h5f is None: + return False + return _fields_present( + calculation, source.data, h5f, open_files, stack, cache, legacy_quantities + ) + + +def _fields_present( + calculation, + data, + h5f, + open_files, + stack, + cache, + legacy_quantities, +): + indices = _valid_indices(data, h5f) + field_names = {field.name for field in dataclasses.fields(data)} + if "valid_indices" in field_names and not indices: + return False + missing = False + for field in dataclasses.fields(data): + if field.name == "valid_indices" or _is_optional(field): + continue + value = getattr(data, field.name) + if check.is_none(value): + continue + if isinstance(value, Link): + if not _confirm_read( + calculation, + value.quantity.lstrip("_"), + value.quantity, + value.source, + open_files, + stack, + cache, + legacy_quantities, + ): + return False + continue + if not _value_present(value, h5f, indices): + missing = True + return None if missing else True + + +def _value_present(value, h5f, indices): + if isinstance(value, Length): + return h5f.get(value.dataset) is not None + path = str(value) + if "{}" in path: + if not indices: + return False + return h5f.get(_format_index(path, indices[0])) is not None + return h5f.get(path) is not None + + +def _valid_indices(data, h5f): + field_names = {field.name for field in dataclasses.fields(data)} + if "valid_indices" not in field_names: + return None + dataset = h5f.get(str(getattr(data, "valid_indices"))) + if dataset is None: + return None + raw_value = dataset[()] + if np.ndim(raw_value) == 0: + return list(range(int(raw_value))) + return [convert.text_to_string(index) for index in raw_value] + + +def _file_version(calculation, open_files, stack): + h5f = _open_h5(calculation, open_files, stack, calculation._file or DEFAULT_FILE) + if h5f is None: + return None + try: + return raw.Version( + int(h5f[schema.version.major][()]), + int(h5f[schema.version.minor][()]), + int(h5f[schema.version.patch][()]), + ) + except (KeyError, OSError, TypeError, ValueError): + return None + + +def _open_h5(calculation, open_files, stack, filename): + if filename in open_files: + return open_files[filename] + try: + handle = stack.enter_context(h5py.File(calculation._path / filename, "r")) + except (FileNotFoundError, OSError): + handle = None + open_files[filename] = handle + return handle + + +def _invoke( + calculation, + call_name, + method_name, + source_name, + legacy_quantities, + convention=None, +): + try: + quantity = _quantity_object(calculation, call_name) + except Exception: + return False + method = getattr(quantity, method_name, None) + if method is None: + return False + if source_name == DEFAULT_SELECTION: + return _call_succeeds(method) + if convention is None: + convention = _source_convention(calculation, call_name, legacy_quantities) + if convention == "keyword": + return _call_succeeds(lambda: method(selection=source_name)) + if convention == "index": + return _call_succeeds(lambda: getattr(quantity[source_name], method_name)()) + return False + + +def _quantity_object(calculation, call_name): + if "." in call_name: + group_name, member = call_name.split(".", 1) + return getattr(getattr(calculation, group_name), member) + return getattr(calculation, call_name) + + +def _is_optional(field): + annotation = field.type + if not isinstance(annotation, str): + annotation = getattr(annotation, "__name__", str(annotation)) + annotation = annotation.strip() + return annotation.startswith("Optional") or annotation.startswith("typing.Optional") + + +def _format_index(path, index): + if isinstance(index, (int, np.integer)): + index = int(index) + 1 + return path.format(index) + + +def _accepts_selection(func): + try: + return "selection" in inspect.signature(func).parameters + except (TypeError, ValueError): + return False + + +def _call_succeeds(func): + try: + func() + return True + except MemoryError: + # Out-of-memory means the data exists but is too large to load. Propagate it + # so the caller can report it differently from a genuinely missing source. + raise + except Exception: + return False + + +def _call_snippet(call_name, method_name, source_name, convention): + access = f"calculation.{call_name}" + if convention == "keyword": + return f"{access}.{method_name}(selection={source_name!r})" + if convention == "index": + return f"{access}[{source_name!r}].{method_name}()" + return f"{access}.{method_name}()" + + +def _out_of_memory_message(call_name, method_name, source_name, convention, error): + snippet = _call_snippet(call_name, method_name, source_name, convention) + reason = str(error).strip() or error.__class__.__name__ + return ( + f"{snippet}: the data is available but could not be loaded because it ran " + f"out of memory ({reason})." + ) diff --git a/src/py4vasp/_util/merge.py b/src/py4vasp/_util/merge.py new file mode 100644 index 000000000..e182523d2 --- /dev/null +++ b/src/py4vasp/_util/merge.py @@ -0,0 +1,77 @@ +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import itertools +from collections.abc import Sequence + +import numpy as np + +from py4vasp import exception + + +def is_unset(value): + if value is None: + return True + if isinstance(value, np.ndarray): + return value.size == 0 + try: + return not value + except ValueError: + return False + + +def values_equal(left_value, right_value): + if isinstance(left_value, np.ndarray) or isinstance(right_value, np.ndarray): + try: + return np.array_equal(np.asarray(left_value), np.asarray(right_value)) + except Exception: + return False + sequence_type = (list, tuple) + if isinstance(left_value, sequence_type) and isinstance(right_value, sequence_type): + if len(left_value) != len(right_value): + return False + return all( + values_equal(left_entry, right_entry) + for left_entry, right_entry in zip(left_value, right_value) + ) + return left_value == right_value + + +def merge_field_or_raise(left_field, right_field, field_name, object_name): + if is_unset(left_field): + return right_field + if is_unset(right_field): + return left_field + if not values_equal(left_field, right_field): + message = f"""Cannot combine two {object_name} with incompatible {field_name}: + left: {left_field} + right: {right_field}""" + raise exception.IncorrectUsage(message) + return left_field + + +def merge_unique_sequences(left_values, right_values, entries_equal): + if is_unset(left_values): + return right_values + if is_unset(right_values): + return left_values + if not isinstance(left_values, Sequence) or not isinstance(right_values, Sequence): + message = "Special merge expected sequence inputs on both sides." + raise exception.IncorrectUsage(message) + + merged = [] + for value in itertools.chain(left_values, right_values): + if any(entries_equal(value, seen) for seen in merged): + continue + merged.append(value) + return _as_left_sequence_type(left_values, merged) + + +def _as_left_sequence_type(left_values, merged): + if isinstance(left_values, list): + return merged + if isinstance(left_values, tuple): + return tuple(merged) + try: + return type(left_values)(merged) + except Exception: + return tuple(merged) diff --git a/tests/calculation/test_band.py b/tests/calculation/test_band.py index 84649a414..63161eac2 100644 --- a/tests/calculation/test_band.py +++ b/tests/calculation/test_band.py @@ -1,747 +1,747 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import types -from dataclasses import fields -from unittest.mock import patch - -import numpy as np -import pytest - -from py4vasp import exception -from py4vasp._calculation.band import _OCCUPATION_CUTOFF, Band, BandHandler -from py4vasp._calculation.kpoint import Kpoint -from py4vasp._calculation.projector import Projector -from py4vasp._raw.data_db import Band_DB -from py4vasp._util import slicing - - -@pytest.fixture -def single_band(raw_data): - raw_band = raw_data.band("single") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.fermi_energy_argument = 1.23 - band.ref.fermi_energy = 0.0 - band.ref.bands = raw_band.dispersion.eigenvalues[0] - band.ref.fermi_energy_argument - band.ref.occupations = raw_band.occupations[0] - band.ref.num_occupied_bands = int( - np.max(np.sum(band.ref.occupations > _OCCUPATION_CUTOFF, axis=-1)) - ) - band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) - band.ref.raw_data = raw_band - return band - - -@pytest.fixture -def multiple_bands(raw_data): - raw_band = raw_data.band("multiple") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.fermi_energy = raw_band.fermi_energy - band.ref.bands = raw_band.dispersion.eigenvalues[0] - raw_band.fermi_energy - band.ref.occupations = raw_band.occupations[0] - band.ref.num_occupied_bands = int( - np.max(np.sum(band.ref.occupations > _OCCUPATION_CUTOFF, axis=-1)) - ) - band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) - band.ref.raw_data = raw_band - return band - - -@pytest.fixture -def with_projectors(raw_data): - raw_band = raw_data.band("multiple with_projectors") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.bands = raw_band.dispersion.eigenvalues[0] - raw_band.fermi_energy - band.ref.Sr = np.sum(raw_band.projections[0, 0:2, :, :, :], axis=(0, 1)) - band.ref.p = np.sum(raw_band.projections[0, :, 1:4, :, :], axis=(0, 1)) - band.ref.selections = Projector.from_data(raw_band.projectors).selections() - return band - - -@pytest.fixture -def line_no_labels(raw_data): - raw_band = raw_data.band("line no_labels") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) - return band - - -@pytest.fixture -def line_with_labels(raw_data): - raw_band = raw_data.band("line with_labels") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) - return band - - -@pytest.fixture -def spin_polarized(raw_data): - raw_band = raw_data.band("spin_polarized") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - assert raw_band.fermi_energy == 0 - band.ref.fermi_energy = raw_band.fermi_energy - band.ref.bands_up = raw_band.dispersion.eigenvalues[0] - band.ref.bands_down = raw_band.dispersion.eigenvalues[1] - band.ref.occupations_up = raw_band.occupations[0] - band.ref.num_occupied_bands_up = int( - np.max(np.sum(band.ref.occupations_up > _OCCUPATION_CUTOFF, axis=-1)) - ) - band.ref.occupations_down = raw_band.occupations[1] - band.ref.num_occupied_bands_down = int( - np.max(np.sum(band.ref.occupations_down > _OCCUPATION_CUTOFF, axis=-1)) - ) - band.ref.raw_data = raw_band - return band - - -@pytest.fixture -def spin_projectors(raw_data): - raw_band = raw_data.band("spin_polarized with_projectors") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.bands_up = raw_band.dispersion.eigenvalues[0] - band.ref.bands_down = raw_band.dispersion.eigenvalues[1] - band.ref.s_up = np.sum(raw_band.projections[0, :, 0, :, :], axis=0) - band.ref.s_down = np.sum(raw_band.projections[1, :, 0, :, :], axis=0) - band.ref.Fe_d_up = np.sum(raw_band.projections[0, 0:3, 2, :, :], axis=0) - band.ref.Fe_d_down = np.sum(raw_band.projections[1, 0:3, 2, :, :], axis=0) - band.ref.O_up = np.sum(raw_band.projections[0, 3:7, :, :, :], axis=(0, 1)) - band.ref.O_down = np.sum(raw_band.projections[1, 3:7, :, :, :], axis=(0, 1)) - projector = Projector.from_data(raw_band.projectors) - band.ref.projectors_string = str(projector) - return band - - -@pytest.fixture -def noncollinear_projectors(raw_data): - raw_band = raw_data.band("noncollinear with_projectors") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.total = np.sum(raw_band.projections[0], axis=(0, 1)) - band.ref.sigma_x = np.sum(raw_band.projections[1], axis=(0, 1)) - band.ref.Ba_sigma_y = np.sum(raw_band.projections[2, 0:2], axis=(0, 1)) - band.ref.d_sigma_z = np.sum(raw_band.projections[3, :, 2], axis=0) - band.ref.occupations = raw_band.occupations[0] - band.ref.num_occupied_bands = int( - np.max(np.sum(band.ref.occupations > _OCCUPATION_CUTOFF, axis=-1)) - ) - band.ref.fermi_energy = raw_band.fermi_energy - band.ref.raw_data = raw_band - return band - - -@pytest.fixture(params=["x~y", "x~z", "y~z"]) -def spin_texture(raw_data, request): - raw_band = raw_data.band(f"spin_texture {request.param}") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - project_all_xy = np.sum(raw_band.projections[1:3, ..., 1], axis=(1, 2)) - project_Pb_xy = np.sum(raw_band.projections[1:3, 2, :, :, 0], axis=1) - project_d_yz = np.sum(raw_band.projections[2:4, :, 2, :, 2], axis=1) - project_5_p_zx = raw_band.projections[[1, 3], 4, 1, :, 0] - cut = {"x~y": "c", "x~z": "b", "y~z": "a"}[request.param] - plot_plane = _spin_texture_plane(raw_band, cut) - band.ref.expected_data = { - "sigma_x~sigma_y_band=2": _expected_quiver_data( - project_all_xy, (0, 1), plot_plane - ), - "Pb_sigma_1~sigma_2_band=1": _expected_quiver_data( - project_Pb_xy, (0, 1), plot_plane - ), - "d_y~z_band=3": _expected_quiver_data(project_d_yz, (1, 2), plot_plane), - "O_2_p_x~z_band=1": _expected_quiver_data(project_5_p_zx, (0, 2), plot_plane), - } - band.ref.expected_lattice = expected_lattice(request.param) - return band - - -def _spin_texture_plane(raw_band, cut): - reciprocal_cell = Kpoint.from_data( - raw_band.dispersion.kpoints - )._reciprocal_lattice_vectors() - return slicing.plane(reciprocal_cell, cut, normal=None) - - -def _expected_quiver_data(two_component_data, axes, plot_plane): - """Embed 2-component spin data into 3D, reshape grid, and project onto plane.""" - nkp1, nkp2 = 4, 3 - embedded = np.zeros( - (3, two_component_data.shape[-1]), dtype=two_component_data.dtype - ) - embedded[list(axes)] = two_component_data - embedded = embedded.reshape(3, nkp2, nkp1).transpose(0, 2, 1) - return slicing._project_vectors_to_plane(plot_plane, embedded) - - -@pytest.fixture -def spin_texture_xy(raw_data): - raw_band = raw_data.band("spin_texture x~y") - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - band.ref.expected_lattice = expected_lattice("x~y") - return band - - -@pytest.fixture -def asymmetric_spin_texture(): - """15x5x1 mesh where sigma_x increases along kx and sigma_y along ky.""" - from py4vasp import _demo, raw - - nkpx, nkpy, nkpz = 15, 5, 1 - num_kpoints = nkpx * nkpy - coordinates = np.array( - [ - (kx, ky, 0.0) - for ky in np.linspace(0, 1, nkpy, endpoint=False) - for kx in np.linspace(0, 1, nkpx, endpoint=False) - ] - ) - cell = raw.Cell(np.diag([5.0, 5.0, 10.0]).astype(float), scale=raw.VaspData(1.0)) - kpoints = raw.Kpoint( - mode="explicit", - number=num_kpoints, - number_x=nkpx, - number_y=nkpy, - number_z=nkpz, - coordinates=coordinates, - weights=np.ones(num_kpoints), - cell=cell, - ) - dispersion = raw.Dispersion(kpoints, np.zeros((4, num_kpoints, 1))) - projectors = _demo.projector.Ba2PbO4(use_orbitals=True) - num_orbitals = len(projectors.orbital_types) - projections = np.zeros((4, _demo.NUMBER_ATOMS, num_orbitals, num_kpoints, 1)) - kx_index = np.tile(np.arange(nkpx), nkpy) - ky_index = np.repeat(np.arange(nkpy), nkpx) - projections[1, :, :, :, 0] = kx_index[None, None, :] - projections[2, :, :, :, 0] = ky_index[None, None, :] - raw_band = raw.Band( - dispersion=dispersion, - fermi_energy=0.0, - occupations=np.ones((4, num_kpoints, 1)), - projectors=projectors, - projections=projections, - ) - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - scale = _demo.NUMBER_ATOMS * num_orbitals - embedded = np.zeros((3, num_kpoints)) - embedded[0] = kx_index.astype(float) * scale - embedded[1] = ky_index.astype(float) * scale - embedded = embedded.reshape(3, nkpy, nkpx).transpose(0, 2, 1) - reciprocal_cell = Kpoint.from_data(kpoints)._reciprocal_lattice_vectors() - plot_plane = slicing.plane(reciprocal_cell, "c", normal=None) - band.ref.expected_data = slicing._project_vectors_to_plane(plot_plane, embedded) - return band - - -@pytest.fixture -def rotated_spin_texture(): - """4x3x1 mesh with a 45-degree rotated cell and uniform sigma_x = 1.""" - from py4vasp import _demo, raw - - nkpx, nkpy, nkpz = 4, 3, 1 - num_kpoints = nkpx * nkpy - coordinates = np.array( - [ - (kx, ky, 0.0) - for ky in np.linspace(0, 1, nkpy, endpoint=False) - for kx in np.linspace(0, 1, nkpx, endpoint=False) - ] - ) - s = 5.0 / np.sqrt(2) - lattice = np.array([[s, s, 0.0], [-s, s, 0.0], [0.0, 0.0, 10.0]]) - cell = raw.Cell(lattice, scale=raw.VaspData(1.0)) - kpoints = raw.Kpoint( - mode="explicit", - number=num_kpoints, - number_x=nkpx, - number_y=nkpy, - number_z=nkpz, - coordinates=coordinates, - weights=np.ones(num_kpoints), - cell=cell, - ) - dispersion = raw.Dispersion(kpoints, np.zeros((4, num_kpoints, 1))) - projectors = _demo.projector.Ba2PbO4(use_orbitals=True) - num_orbitals = len(projectors.orbital_types) - projections = np.zeros((4, _demo.NUMBER_ATOMS, num_orbitals, num_kpoints, 1)) - projections[1, :, :, :, 0] = 1.0 # uniform sigma_x - raw_band = raw.Band( - dispersion=dispersion, - fermi_energy=0.0, - occupations=np.ones((4, num_kpoints, 1)), - projectors=projectors, - projections=projections, - ) - band = Band.from_data(raw_band) - band.ref = types.SimpleNamespace() - scale = _demo.NUMBER_ATOMS * num_orbitals - embedded = np.zeros((3, num_kpoints)) - embedded[0] = scale - embedded = embedded.reshape(3, nkpy, nkpx).transpose(0, 2, 1) - reciprocal_cell = Kpoint.from_data(kpoints)._reciprocal_lattice_vectors() - plot_plane = slicing.plane(reciprocal_cell, "c", normal=None) - band.ref.expected_data = slicing._project_vectors_to_plane(plot_plane, embedded) - return band - - -def expected_lattice(selection): - if selection == "x~y": - return np.array([[1.52216787, 0.0], [0.14521927, 1.51522486]]) - else: - return np.array([[1.52216787, 0.0], [0.29043854, 0.89433656]]) - - -def test_single_band_read(single_band, Assert): - band = single_band.read(fermi_energy=single_band.ref.fermi_energy_argument) - assert band["fermi_energy"] == single_band.ref.fermi_energy - Assert.allclose(band["bands"], single_band.ref.bands) - Assert.allclose(band["occupations"], single_band.ref.occupations) - Assert.allclose(band["kpoint_distances"], single_band.ref.kpoints.distances()) - assert "kpoint_labels" not in band - assert "projections" not in band - - -def test_multiple_bands_read(multiple_bands, Assert): - band = multiple_bands.read() - assert band["fermi_energy"] == multiple_bands.ref.fermi_energy - Assert.allclose(band["bands"], multiple_bands.ref.bands) - Assert.allclose(band["occupations"], multiple_bands.ref.occupations) - - -def test_with_projectors_read(with_projectors, Assert): - band = with_projectors.read("Sr p") - Assert.allclose(band["Sr"], with_projectors.ref.Sr) - Assert.allclose(band["p"], with_projectors.ref.p) - - -def test_line_with_labels_read(line_with_labels, Assert): - band = line_with_labels.read() - Assert.allclose(band["kpoint_distances"], line_with_labels.ref.kpoints.distances()) - assert band["kpoint_labels"] == line_with_labels.ref.kpoints.labels() - - -def test_spin_polarized_read(spin_polarized, Assert): - band = spin_polarized.read() - Assert.allclose(band["bands_up"], spin_polarized.ref.bands_up) - Assert.allclose(band["bands_down"], spin_polarized.ref.bands_down) - Assert.allclose(band["occupations_up"], spin_polarized.ref.occupations_up) - Assert.allclose(band["occupations_down"], spin_polarized.ref.occupations_down) - - -def test_spin_projectors_read(spin_projectors, Assert): - band = spin_projectors.read(selection="s Fe(d)") - Assert.allclose(band["s_up"], spin_projectors.ref.s_up) - Assert.allclose(band["s_down"], spin_projectors.ref.s_down) - Assert.allclose(band["Fe_d_up"], spin_projectors.ref.Fe_d_up) - Assert.allclose(band["Fe_d_down"], spin_projectors.ref.Fe_d_down) - - -def test_noncollinear_projectors_read(noncollinear_projectors, Assert): - band = noncollinear_projectors.read(selection="total sigma_x Ba(y) d(sigma_3)") - Assert.allclose(band["total"], noncollinear_projectors.ref.total) - Assert.allclose(band["sigma_x"], noncollinear_projectors.ref.sigma_x) - Assert.allclose(band["Ba_y"], noncollinear_projectors.ref.Ba_sigma_y) - Assert.allclose(band["d_sigma_3"], noncollinear_projectors.ref.d_sigma_z) - - -def test_combining_projections(with_projectors, Assert): - band = with_projectors.read("Sr + p, Sr - p") - addition = with_projectors.ref.Sr + with_projectors.ref.p - subtraction = with_projectors.ref.Sr - with_projectors.ref.p - Assert.allclose(band["Sr + p"], addition) - Assert.allclose(band["Sr - p"], subtraction) - - -def test_more_projections_style(raw_data, Assert): - """Vasp 6.1 may store more orbital types then projections available. This - test checks that this does not lead to any issues when an available element - is used.""" - raw_band = raw_data.band("spin_polarized excess_orbitals") - band = Band.from_data(raw_band).read("Fe g") - zero = np.zeros_like(band["Fe_up"]) - Assert.allclose(band["g_up"], zero) - Assert.allclose(band["g_down"], zero) - - -def test_single_polarized_to_frame(single_band, Assert, not_core): - actual = single_band.to_frame(fermi_energy=single_band.ref.fermi_energy_argument) - Assert.allclose(actual.bands, single_band.ref.bands[:, 0]) - Assert.allclose(actual.occupations, single_band.ref.occupations[:, 0]) - Assert.allclose(actual.kpoint_distances, single_band.ref.kpoints.distances()) - - -def test_multiple_bands_to_frame(multiple_bands, Assert, not_core): - actual = multiple_bands.to_frame() - Assert.allclose(actual.bands, multiple_bands.ref.bands.T.flatten()) - Assert.allclose(actual.occupations, multiple_bands.ref.occupations.T.flatten()) - kpoint_distances = np.repeat(multiple_bands.ref.kpoints.distances(), repeats=3) - Assert.allclose(actual.kpoint_distances, kpoint_distances) - - -def test_line_with_labels_to_frame(line_with_labels, Assert, not_core): - actual = line_with_labels.to_frame() - kpoint_distances = np.repeat(line_with_labels.ref.kpoints.distances(), repeats=3) - kpoint_labels = np.repeat(line_with_labels.ref.kpoints.labels(), repeats=3) - actual_kpoint_labels = np.array(actual.kpoint_labels).astype(np.str_) - Assert.allclose(actual.kpoint_distances, kpoint_distances) - Assert.allclose(actual_kpoint_labels, kpoint_labels) - - -def test_with_projectors_to_frame(with_projectors, Assert, not_core): - actual = with_projectors.to_frame("Sr p") - Assert.allclose(actual.Sr, with_projectors.ref.Sr.T.flatten()) - Assert.allclose(actual.p, with_projectors.ref.p.T.flatten()) - - -def test_spin_polarized_to_frame(spin_polarized, Assert, not_core): - actual = spin_polarized.to_frame() - ref = spin_polarized.ref - Assert.allclose(actual.bands_up, ref.bands_up.T.flatten()) - Assert.allclose(actual.bands_down, ref.bands_down.T.flatten()) - Assert.allclose(actual.occupations_up, ref.occupations_up.T.flatten()) - Assert.allclose(actual.occupations_down, ref.occupations_down.T.flatten()) - - -def test_spin_projectors_to_frame(spin_projectors, Assert, not_core): - actual = spin_projectors.to_frame(selection="O Fe(d)") - Assert.allclose(actual.O_up, spin_projectors.ref.O_up.T.flatten()) - Assert.allclose(actual.O_down, spin_projectors.ref.O_down.T.flatten()) - Assert.allclose(actual.Fe_d_up, spin_projectors.ref.Fe_d_up.T.flatten()) - Assert.allclose(actual.Fe_d_down, spin_projectors.ref.Fe_d_down.T.flatten()) - - -def test_single_band_plot(single_band, Assert): - fig = single_band.plot(fermi_energy=single_band.ref.fermi_energy_argument) - assert fig.ylabel == "Energy (eV)" - assert len(fig.series) == 1 - assert fig.series[0].weight is None - Assert.allclose(fig.series[0].x, single_band.ref.kpoints.distances()) - Assert.allclose(fig.series[0].y, single_band.ref.bands.T) - - -def test_multiple_bands_plot(multiple_bands, Assert): - fig = multiple_bands.plot() - assert len(fig.series) == 1 # all bands in one plot - assert len(fig.series[0].x) == fig.series[0].y.shape[-1] - Assert.allclose(fig.series[0].y, multiple_bands.ref.bands.T) - - -def test_with_projectors_plot_default_width(with_projectors, Assert): - default_width = 0.5 - fig = with_projectors.plot(selection="Sr, p") - check_figure(fig, default_width, with_projectors.ref, Assert) - - -def test_with_projectors_plot_custom_width(with_projectors, Assert): - width = 0.1 - fig = with_projectors.plot(selection="Sr, p", width=width) - check_figure(fig, width, with_projectors.ref, Assert) - - -def test_spin_projectors_plot(spin_projectors, Assert): - reference = spin_projectors.ref - width = 0.05 - fig = spin_projectors.plot("O", width=width) - assert len(fig.series) == 2 - assert fig.series[0].label == "O_up" - check_data(fig.series[0], width, reference.bands_up, reference.O_up, Assert) - assert fig.series[1].label == "O_down" - check_data(fig.series[1], width, reference.bands_down, reference.O_down, Assert) - - -def check_figure(fig, weight, reference, Assert): - assert len(fig.series) == 2 - assert fig.series[0].label == "Sr" - assert fig.series[1].label == "p" - check_data(fig.series[0], weight, reference.bands, reference.Sr, Assert) - check_data(fig.series[1], weight, reference.bands, reference.p, Assert) - - -def check_data(series, weight, band, projection, Assert): - assert len(series.x) == series.y.shape[-1] - assert series.y.shape == series.weight.shape - Assert.allclose(series.y, band.T) - Assert.allclose(series.weight, weight * projection.T) - - -def test_spin_polarized_plot(spin_polarized, Assert): - fig = spin_polarized.plot() - assert len(fig.series) == 2 - assert fig.series[0].label == "up" - Assert.allclose(fig.series[0].y, spin_polarized.ref.bands_up.T) - assert fig.series[1].label == "down" - Assert.allclose(fig.series[1].y, spin_polarized.ref.bands_down.T) - - -def test_line_no_labels_plot(line_no_labels, Assert): - fig = line_no_labels.plot() - check_ticks(fig, line_no_labels.ref.kpoints, Assert) - reference_labels = ( - "$[0 0 0]$", - "$[0 0 \\frac{1}{2}]$", - "$[\\frac{1}{2} \\frac{1}{2} \\frac{1}{2}]$|$[0 0 0]$", - "$[\\frac{1}{2} \\frac{1}{2} 0]$", - "$[\\frac{1}{2} \\frac{1}{2} \\frac{1}{2}]$", - ) - assert tuple(fig.xticks.values()) == reference_labels - - -def test_line_with_labels_plot(line_with_labels, Assert): - fig = line_with_labels.plot() - check_ticks(fig, line_with_labels.ref.kpoints, Assert) - assert tuple(fig.xticks.values()) == (r"$\Gamma$", "", r"M|$\Gamma$", "Y", "M") - - -def test_noncollinear_plot(noncollinear_projectors, Assert): - default_width = 0.5 - fig = noncollinear_projectors.plot("total sigma_x") - assert len(fig.series) == 2 - assert fig.series[0].label == "total" - reference = noncollinear_projectors.ref - Assert.allclose(fig.series[0].weight, default_width * reference.total.T) - assert fig.series[0].marker is None - assert fig.series[1].label == "sigma_x" - Assert.allclose(fig.series[1].weight, reference.sigma_x.T) - assert fig.series[1].marker == "o" - assert fig.series[1].weight_mode == "color" - - -def check_ticks(fig, kpoints, Assert): - dists = kpoints.distances() - xticks = (*dists[:: kpoints.line_length()], dists[-1]) - Assert.allclose(list(fig.xticks.keys()), np.array(xticks)) - - -def test_plot_incorrect_width(with_projectors): - with pytest.raises(exception.IncorrectUsage): - with_projectors.plot("Sr", width="not a number") - - -@patch.object(Band, "to_graph") -def test_to_plotly(mock_plot, single_band): - fig = single_band.to_plotly("selection", width=0.2) - mock_plot.assert_called_once_with("selection", width=0.2) - graph = mock_plot.return_value - graph.to_plotly.assert_called_once() - assert fig == graph.to_plotly.return_value - - -def test_to_image(single_band): - check_to_image(single_band, None, "band.png") - custom_filename = "custom.jpg" - check_to_image(single_band, custom_filename, custom_filename) - - -def check_to_image(single_band, filename_argument, expected_filename): - with patch.object(Band, "to_plotly") as plot: - single_band.to_image("args", filename=filename_argument, key="word") - plot.assert_called_once_with("args", key="word") - fig = plot.return_value - fig.write_image.assert_called_once_with(single_band._path / expected_filename) - - -def test_band_selections(with_projectors): - actual = with_projectors.selections() - actual.pop("band") # remove band selections - assert actual == with_projectors.ref.selections - - -def test_to_quiver_with_incorrect_selection_raises_error(spin_texture_xy): - with pytest.raises(exception.IncorrectUsage): - spin_texture_xy.to_quiver("x") - with pytest.raises(exception.IncorrectUsage): - spin_texture_xy.to_quiver("x~y") - with pytest.raises(exception.IncorrectUsage): - spin_texture_xy.to_quiver("band=2") - - -@pytest.mark.parametrize( - "selection", - [ - "band=2(sigma_x~sigma_y)", - "Pb(sigma_1~sigma_2(band=1))", - "band=3(d(y~z))", - "p(5(band=1(z~x)))", - ], -) -def test_band_to_quiver(spin_texture, selection, Assert): - graph = spin_texture.to_quiver(selection) - assert graph.title == "Spin Texture" - assert len(graph) == 1 - series = graph.series[0] - Assert.allclose( - series.lattice.vectors, spin_texture.ref.expected_lattice, tolerance=1e6 - ) - assert series.label in spin_texture.ref.expected_data - Assert.allclose(series.data, spin_texture.ref.expected_data[series.label]) - - -@pytest.mark.parametrize("normal", [None, "x", "y", "z"]) -def test_band_to_quiver_normal(spin_texture_xy, normal, Assert): - graph = spin_texture_xy.to_quiver("band=1(x~z)", normal=normal) - assert len(graph) == 1 - quiver_plot = graph.series[0] - if normal == "x": - expected_lattice = [[-0.32989453, 1.48598944], [1.44773855, 0.47015754]] - elif normal == "y": - expected_lattice = [[1.44773855, 0.47015754], [-0.32989453, 1.48598944]] - elif normal == "z": - expected_lattice = [[1.52043113, 0.07269258], [0.07269258, 1.52043113]] - else: - expected_lattice = spin_texture_xy.ref.expected_lattice - Assert.allclose(quiver_plot.lattice.vectors, expected_lattice, tolerance=1e6) - - -@pytest.mark.parametrize( - "supercell, expected_supercell", [(None, (1, 1)), (2, (2, 2)), ([2, 4], (2, 4))] -) -def test_band_to_quiver_supercell( - spin_texture_xy, supercell, expected_supercell, Assert -): - graph = spin_texture_xy.to_quiver("band=3(sigma_2~sigma_1)", supercell=supercell) - assert len(graph) == 1 - quiver_plot = graph.series[0] - Assert.allclose(quiver_plot.supercell, expected_supercell) - - -def test_band_to_quiver_asymmetric_mesh(asymmetric_spin_texture, Assert): - """With a heavily asymmetric k-mesh (15x5x1), verify correct axis ordering.""" - graph = asymmetric_spin_texture.to_quiver("x~y(band=1)") - quiver_data = graph.series[0].data - assert quiver_data.shape[1] == 15 - assert quiver_data.shape[2] == 5 - Assert.allclose(quiver_data, asymmetric_spin_texture.ref.expected_data) - - -def test_band_to_quiver_rotation_projects_spin_vectors(rotated_spin_texture, Assert): - """Verify that spin vectors are rotated into the plot frame.""" - graph = rotated_spin_texture.to_quiver("x~y(band=1)") - quiver_data = graph.series[0].data - assert quiver_data.shape == (2, 4, 3) - # A pure Cartesian x-vector should project onto both plot axes for a rotated cell - assert not np.allclose(quiver_data[1], 0), ( - "Rotation is not applied: sigma_x should project onto both plot axes " - "for a 45-degree rotated cell" - ) - Assert.allclose(quiver_data, rotated_spin_texture.ref.expected_data) - - -def test_multiple_bands_print(multiple_bands, format_): - actual, _ = format_(multiple_bands) - reference = f""" -band data: - 48 k-points - 3 bands -no projectors - """.strip() - assert actual == {"text/plain": reference} - - -def test_line_no_labels_print(line_no_labels, format_): - actual, _ = format_(line_no_labels) - reference = f""" -band data: - 20 k-points - 3 bands -no projectors - """.strip() - assert actual == {"text/plain": reference} - - -def test_line_with_labels_print(line_with_labels, format_): - actual, _ = format_(line_with_labels) - reference = f""" -band data: - 20 k-points - 3 bands -no projectors - """.strip() - assert actual == {"text/plain": reference} - - -def test_spin_projectors_print(spin_projectors, format_): - actual, _ = format_(spin_projectors) - reference = f""" -spin polarized band data: - 48 k-points - 3 bands -{spin_projectors.ref.projectors_string} - """.strip() - assert actual == {"text/plain": reference} - - -def _check_to_database(_band): - handler = BandHandler.from_data(_band.ref.raw_data) - database_data: Band_DB = handler.to_database( - fermi_energy=getattr(_band.ref, "fermi_energy_argument", None) - ) - - assert isinstance(database_data, Band_DB) - - assert database_data.fermi_energy_raw == _band.ref.fermi_energy - assert database_data.fermi_energy == getattr( - _band.ref, "fermi_energy_argument", _band.ref.fermi_energy - ) - - if getattr(_band.ref, "num_occupied_bands", None) is not None: - assert database_data.num_occupied_bands == _band.ref.num_occupied_bands - elif ( - getattr(_band.ref, "occupations_up", None) is not None - and getattr(_band.ref, "occupations_down", None) is not None - ): - assert database_data.num_occupied_bands_up == _band.ref.num_occupied_bands_up - assert ( - database_data.num_occupied_bands_down == _band.ref.num_occupied_bands_down - ) - - for fld in fields(Band_DB): - if fld.name.startswith("num"): - assert getattr(database_data, fld.name) is None or isinstance( - getattr(database_data, fld.name), int - ), f"{fld.name} has unexpected type {type(getattr(database_data, fld.name))}: {getattr(database_data, fld.name)}" - else: - assert ( - getattr(database_data, fld.name) is None - or isinstance(getattr(database_data, fld.name), float) - or ( - fld.name.startswith("__") - and isinstance(getattr(database_data, fld.name), str) - ) - ), f"{fld.name} has unexpected type {type(getattr(database_data, fld.name))}: {getattr(database_data, fld.name)}" - - -def test_to_database_single_band(single_band): - _check_to_database(single_band) - - -def test_to_database_multiple_bands(multiple_bands): - _check_to_database(multiple_bands) - - -def test_to_database_spin_polarized(spin_polarized): - _check_to_database(spin_polarized) - - -def test_to_database_noncollinear_projectors(noncollinear_projectors): - _check_to_database(noncollinear_projectors) - - -def test_dispatcher_to_database_default(single_band): - """Dispatcher._to_database() must return {selection_name: handler_result}.""" - result = single_band._to_database() - assert isinstance(result, dict) - assert "band" in result - assert isinstance(result["band"], Band_DB) - - -def test_factory_methods(raw_data, check_factory_methods): - data = raw_data.band("multiple") - parameters = {"to_quiver": {"selection": "x~y(band=1)"}} - check_factory_methods(Band, data, parameters) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import types +from dataclasses import fields +from unittest.mock import patch + +import numpy as np +import pytest + +from py4vasp import exception +from py4vasp._calculation.band import _OCCUPATION_CUTOFF, Band, BandHandler +from py4vasp._calculation.kpoint import Kpoint +from py4vasp._calculation.projector import Projector +from py4vasp._raw.data_db import Band_DB +from py4vasp._util import slicing + + +@pytest.fixture +def single_band(raw_data): + raw_band = raw_data.band("single") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.fermi_energy_argument = 1.23 + band.ref.fermi_energy = 0.0 + band.ref.bands = raw_band.dispersion.eigenvalues[0] - band.ref.fermi_energy_argument + band.ref.occupations = raw_band.occupations[0] + band.ref.num_occupied_bands = int( + np.max(np.sum(band.ref.occupations > _OCCUPATION_CUTOFF, axis=-1)) + ) + band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) + band.ref.raw_data = raw_band + return band + + +@pytest.fixture +def multiple_bands(raw_data): + raw_band = raw_data.band("multiple") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.fermi_energy = raw_band.fermi_energy + band.ref.bands = raw_band.dispersion.eigenvalues[0] - raw_band.fermi_energy + band.ref.occupations = raw_band.occupations[0] + band.ref.num_occupied_bands = int( + np.max(np.sum(band.ref.occupations > _OCCUPATION_CUTOFF, axis=-1)) + ) + band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) + band.ref.raw_data = raw_band + return band + + +@pytest.fixture +def with_projectors(raw_data): + raw_band = raw_data.band("multiple with_projectors") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.bands = raw_band.dispersion.eigenvalues[0] - raw_band.fermi_energy + band.ref.Sr = np.sum(raw_band.projections[0, 0:2, :, :, :], axis=(0, 1)) + band.ref.p = np.sum(raw_band.projections[0, :, 1:4, :, :], axis=(0, 1)) + band.ref.selections = Projector.from_data(raw_band.projectors).selections() + return band + + +@pytest.fixture +def line_no_labels(raw_data): + raw_band = raw_data.band("line no_labels") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) + return band + + +@pytest.fixture +def line_with_labels(raw_data): + raw_band = raw_data.band("line with_labels") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.kpoints = Kpoint.from_data(raw_band.dispersion.kpoints) + return band + + +@pytest.fixture +def spin_polarized(raw_data): + raw_band = raw_data.band("spin_polarized") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + assert raw_band.fermi_energy == 0 + band.ref.fermi_energy = raw_band.fermi_energy + band.ref.bands_up = raw_band.dispersion.eigenvalues[0] + band.ref.bands_down = raw_band.dispersion.eigenvalues[1] + band.ref.occupations_up = raw_band.occupations[0] + band.ref.num_occupied_bands_up = int( + np.max(np.sum(band.ref.occupations_up > _OCCUPATION_CUTOFF, axis=-1)) + ) + band.ref.occupations_down = raw_band.occupations[1] + band.ref.num_occupied_bands_down = int( + np.max(np.sum(band.ref.occupations_down > _OCCUPATION_CUTOFF, axis=-1)) + ) + band.ref.raw_data = raw_band + return band + + +@pytest.fixture +def spin_projectors(raw_data): + raw_band = raw_data.band("spin_polarized with_projectors") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.bands_up = raw_band.dispersion.eigenvalues[0] + band.ref.bands_down = raw_band.dispersion.eigenvalues[1] + band.ref.s_up = np.sum(raw_band.projections[0, :, 0, :, :], axis=0) + band.ref.s_down = np.sum(raw_band.projections[1, :, 0, :, :], axis=0) + band.ref.Fe_d_up = np.sum(raw_band.projections[0, 0:3, 2, :, :], axis=0) + band.ref.Fe_d_down = np.sum(raw_band.projections[1, 0:3, 2, :, :], axis=0) + band.ref.O_up = np.sum(raw_band.projections[0, 3:7, :, :, :], axis=(0, 1)) + band.ref.O_down = np.sum(raw_band.projections[1, 3:7, :, :, :], axis=(0, 1)) + projector = Projector.from_data(raw_band.projectors) + band.ref.projectors_string = str(projector) + return band + + +@pytest.fixture +def noncollinear_projectors(raw_data): + raw_band = raw_data.band("noncollinear with_projectors") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.total = np.sum(raw_band.projections[0], axis=(0, 1)) + band.ref.sigma_x = np.sum(raw_band.projections[1], axis=(0, 1)) + band.ref.Ba_sigma_y = np.sum(raw_band.projections[2, 0:2], axis=(0, 1)) + band.ref.d_sigma_z = np.sum(raw_band.projections[3, :, 2], axis=0) + band.ref.occupations = raw_band.occupations[0] + band.ref.num_occupied_bands = int( + np.max(np.sum(band.ref.occupations > _OCCUPATION_CUTOFF, axis=-1)) + ) + band.ref.fermi_energy = raw_band.fermi_energy + band.ref.raw_data = raw_band + return band + + +@pytest.fixture(params=["x~y", "x~z", "y~z"]) +def spin_texture(raw_data, request): + raw_band = raw_data.band(f"spin_texture {request.param}") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + project_all_xy = np.sum(raw_band.projections[1:3, ..., 1], axis=(1, 2)) + project_Pb_xy = np.sum(raw_band.projections[1:3, 2, :, :, 0], axis=1) + project_d_yz = np.sum(raw_band.projections[2:4, :, 2, :, 2], axis=1) + project_5_p_zx = raw_band.projections[[1, 3], 4, 1, :, 0] + cut = {"x~y": "c", "x~z": "b", "y~z": "a"}[request.param] + plot_plane = _spin_texture_plane(raw_band, cut) + band.ref.expected_data = { + "sigma_x~sigma_y_band=2": _expected_quiver_data( + project_all_xy, (0, 1), plot_plane + ), + "Pb_sigma_1~sigma_2_band=1": _expected_quiver_data( + project_Pb_xy, (0, 1), plot_plane + ), + "d_y~z_band=3": _expected_quiver_data(project_d_yz, (1, 2), plot_plane), + "O_2_p_x~z_band=1": _expected_quiver_data(project_5_p_zx, (0, 2), plot_plane), + } + band.ref.expected_lattice = expected_lattice(request.param) + return band + + +def _spin_texture_plane(raw_band, cut): + reciprocal_cell = Kpoint.from_data( + raw_band.dispersion.kpoints + )._reciprocal_lattice_vectors() + return slicing.plane(reciprocal_cell, cut, normal=None) + + +def _expected_quiver_data(two_component_data, axes, plot_plane): + """Embed 2-component spin data into 3D, reshape grid, and project onto plane.""" + nkp1, nkp2 = 4, 3 + embedded = np.zeros( + (3, two_component_data.shape[-1]), dtype=two_component_data.dtype + ) + embedded[list(axes)] = two_component_data + embedded = embedded.reshape(3, nkp2, nkp1).transpose(0, 2, 1) + return slicing._project_vectors_to_plane(plot_plane, embedded) + + +@pytest.fixture +def spin_texture_xy(raw_data): + raw_band = raw_data.band("spin_texture x~y") + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + band.ref.expected_lattice = expected_lattice("x~y") + return band + + +@pytest.fixture +def asymmetric_spin_texture(): + """15x5x1 mesh where sigma_x increases along kx and sigma_y along ky.""" + from py4vasp import _demo, raw + + nkpx, nkpy, nkpz = 15, 5, 1 + num_kpoints = nkpx * nkpy + coordinates = np.array( + [ + (kx, ky, 0.0) + for ky in np.linspace(0, 1, nkpy, endpoint=False) + for kx in np.linspace(0, 1, nkpx, endpoint=False) + ] + ) + cell = raw.Cell(np.diag([5.0, 5.0, 10.0]).astype(float), scale=raw.VaspData(1.0)) + kpoints = raw.Kpoint( + mode="explicit", + number=num_kpoints, + number_x=nkpx, + number_y=nkpy, + number_z=nkpz, + coordinates=coordinates, + weights=np.ones(num_kpoints), + cell=cell, + ) + dispersion = raw.Dispersion(kpoints, np.zeros((4, num_kpoints, 1))) + projectors = _demo.projector.Ba2PbO4(use_orbitals=True) + num_orbitals = len(projectors.orbital_types) + projections = np.zeros((4, _demo.NUMBER_ATOMS, num_orbitals, num_kpoints, 1)) + kx_index = np.tile(np.arange(nkpx), nkpy) + ky_index = np.repeat(np.arange(nkpy), nkpx) + projections[1, :, :, :, 0] = kx_index[None, None, :] + projections[2, :, :, :, 0] = ky_index[None, None, :] + raw_band = raw.Band( + dispersion=dispersion, + fermi_energy=0.0, + occupations=np.ones((4, num_kpoints, 1)), + projectors=projectors, + projections=projections, + ) + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + scale = _demo.NUMBER_ATOMS * num_orbitals + embedded = np.zeros((3, num_kpoints)) + embedded[0] = kx_index.astype(float) * scale + embedded[1] = ky_index.astype(float) * scale + embedded = embedded.reshape(3, nkpy, nkpx).transpose(0, 2, 1) + reciprocal_cell = Kpoint.from_data(kpoints)._reciprocal_lattice_vectors() + plot_plane = slicing.plane(reciprocal_cell, "c", normal=None) + band.ref.expected_data = slicing._project_vectors_to_plane(plot_plane, embedded) + return band + + +@pytest.fixture +def rotated_spin_texture(): + """4x3x1 mesh with a 45-degree rotated cell and uniform sigma_x = 1.""" + from py4vasp import _demo, raw + + nkpx, nkpy, nkpz = 4, 3, 1 + num_kpoints = nkpx * nkpy + coordinates = np.array( + [ + (kx, ky, 0.0) + for ky in np.linspace(0, 1, nkpy, endpoint=False) + for kx in np.linspace(0, 1, nkpx, endpoint=False) + ] + ) + s = 5.0 / np.sqrt(2) + lattice = np.array([[s, s, 0.0], [-s, s, 0.0], [0.0, 0.0, 10.0]]) + cell = raw.Cell(lattice, scale=raw.VaspData(1.0)) + kpoints = raw.Kpoint( + mode="explicit", + number=num_kpoints, + number_x=nkpx, + number_y=nkpy, + number_z=nkpz, + coordinates=coordinates, + weights=np.ones(num_kpoints), + cell=cell, + ) + dispersion = raw.Dispersion(kpoints, np.zeros((4, num_kpoints, 1))) + projectors = _demo.projector.Ba2PbO4(use_orbitals=True) + num_orbitals = len(projectors.orbital_types) + projections = np.zeros((4, _demo.NUMBER_ATOMS, num_orbitals, num_kpoints, 1)) + projections[1, :, :, :, 0] = 1.0 # uniform sigma_x + raw_band = raw.Band( + dispersion=dispersion, + fermi_energy=0.0, + occupations=np.ones((4, num_kpoints, 1)), + projectors=projectors, + projections=projections, + ) + band = Band.from_data(raw_band) + band.ref = types.SimpleNamespace() + scale = _demo.NUMBER_ATOMS * num_orbitals + embedded = np.zeros((3, num_kpoints)) + embedded[0] = scale + embedded = embedded.reshape(3, nkpy, nkpx).transpose(0, 2, 1) + reciprocal_cell = Kpoint.from_data(kpoints)._reciprocal_lattice_vectors() + plot_plane = slicing.plane(reciprocal_cell, "c", normal=None) + band.ref.expected_data = slicing._project_vectors_to_plane(plot_plane, embedded) + return band + + +def expected_lattice(selection): + if selection == "x~y": + return np.array([[1.52216787, 0.0], [0.14521927, 1.51522486]]) + else: + return np.array([[1.52216787, 0.0], [0.29043854, 0.89433656]]) + + +def test_single_band_read(single_band, Assert): + band = single_band.read(fermi_energy=single_band.ref.fermi_energy_argument) + assert band["fermi_energy"] == single_band.ref.fermi_energy + Assert.allclose(band["bands"], single_band.ref.bands) + Assert.allclose(band["occupations"], single_band.ref.occupations) + Assert.allclose(band["kpoint_distances"], single_band.ref.kpoints.distances()) + assert "kpoint_labels" not in band + assert "projections" not in band + + +def test_multiple_bands_read(multiple_bands, Assert): + band = multiple_bands.read() + assert band["fermi_energy"] == multiple_bands.ref.fermi_energy + Assert.allclose(band["bands"], multiple_bands.ref.bands) + Assert.allclose(band["occupations"], multiple_bands.ref.occupations) + + +def test_with_projectors_read(with_projectors, Assert): + band = with_projectors.read("Sr p") + Assert.allclose(band["Sr"], with_projectors.ref.Sr) + Assert.allclose(band["p"], with_projectors.ref.p) + + +def test_line_with_labels_read(line_with_labels, Assert): + band = line_with_labels.read() + Assert.allclose(band["kpoint_distances"], line_with_labels.ref.kpoints.distances()) + assert band["kpoint_labels"] == line_with_labels.ref.kpoints.labels() + + +def test_spin_polarized_read(spin_polarized, Assert): + band = spin_polarized.read() + Assert.allclose(band["bands_up"], spin_polarized.ref.bands_up) + Assert.allclose(band["bands_down"], spin_polarized.ref.bands_down) + Assert.allclose(band["occupations_up"], spin_polarized.ref.occupations_up) + Assert.allclose(band["occupations_down"], spin_polarized.ref.occupations_down) + + +def test_spin_projectors_read(spin_projectors, Assert): + band = spin_projectors.read(selection="s Fe(d)") + Assert.allclose(band["s_up"], spin_projectors.ref.s_up) + Assert.allclose(band["s_down"], spin_projectors.ref.s_down) + Assert.allclose(band["Fe_d_up"], spin_projectors.ref.Fe_d_up) + Assert.allclose(band["Fe_d_down"], spin_projectors.ref.Fe_d_down) + + +def test_noncollinear_projectors_read(noncollinear_projectors, Assert): + band = noncollinear_projectors.read(selection="total sigma_x Ba(y) d(sigma_3)") + Assert.allclose(band["total"], noncollinear_projectors.ref.total) + Assert.allclose(band["sigma_x"], noncollinear_projectors.ref.sigma_x) + Assert.allclose(band["Ba_y"], noncollinear_projectors.ref.Ba_sigma_y) + Assert.allclose(band["d_sigma_3"], noncollinear_projectors.ref.d_sigma_z) + + +def test_combining_projections(with_projectors, Assert): + band = with_projectors.read("Sr + p, Sr - p") + addition = with_projectors.ref.Sr + with_projectors.ref.p + subtraction = with_projectors.ref.Sr - with_projectors.ref.p + Assert.allclose(band["Sr + p"], addition) + Assert.allclose(band["Sr - p"], subtraction) + + +def test_more_projections_style(raw_data, Assert): + """Vasp 6.1 may store more orbital types then projections available. This + test checks that this does not lead to any issues when an available element + is used.""" + raw_band = raw_data.band("spin_polarized excess_orbitals") + band = Band.from_data(raw_band).read("Fe g") + zero = np.zeros_like(band["Fe_up"]) + Assert.allclose(band["g_up"], zero) + Assert.allclose(band["g_down"], zero) + + +def test_single_polarized_to_frame(single_band, Assert, not_core): + actual = single_band.to_frame(fermi_energy=single_band.ref.fermi_energy_argument) + Assert.allclose(actual.bands, single_band.ref.bands[:, 0]) + Assert.allclose(actual.occupations, single_band.ref.occupations[:, 0]) + Assert.allclose(actual.kpoint_distances, single_band.ref.kpoints.distances()) + + +def test_multiple_bands_to_frame(multiple_bands, Assert, not_core): + actual = multiple_bands.to_frame() + Assert.allclose(actual.bands, multiple_bands.ref.bands.T.flatten()) + Assert.allclose(actual.occupations, multiple_bands.ref.occupations.T.flatten()) + kpoint_distances = np.repeat(multiple_bands.ref.kpoints.distances(), repeats=3) + Assert.allclose(actual.kpoint_distances, kpoint_distances) + + +def test_line_with_labels_to_frame(line_with_labels, Assert, not_core): + actual = line_with_labels.to_frame() + kpoint_distances = np.repeat(line_with_labels.ref.kpoints.distances(), repeats=3) + kpoint_labels = np.repeat(line_with_labels.ref.kpoints.labels(), repeats=3) + actual_kpoint_labels = np.array(actual.kpoint_labels).astype(np.str_) + Assert.allclose(actual.kpoint_distances, kpoint_distances) + Assert.allclose(actual_kpoint_labels, kpoint_labels) + + +def test_with_projectors_to_frame(with_projectors, Assert, not_core): + actual = with_projectors.to_frame("Sr p") + Assert.allclose(actual.Sr, with_projectors.ref.Sr.T.flatten()) + Assert.allclose(actual.p, with_projectors.ref.p.T.flatten()) + + +def test_spin_polarized_to_frame(spin_polarized, Assert, not_core): + actual = spin_polarized.to_frame() + ref = spin_polarized.ref + Assert.allclose(actual.bands_up, ref.bands_up.T.flatten()) + Assert.allclose(actual.bands_down, ref.bands_down.T.flatten()) + Assert.allclose(actual.occupations_up, ref.occupations_up.T.flatten()) + Assert.allclose(actual.occupations_down, ref.occupations_down.T.flatten()) + + +def test_spin_projectors_to_frame(spin_projectors, Assert, not_core): + actual = spin_projectors.to_frame(selection="O Fe(d)") + Assert.allclose(actual.O_up, spin_projectors.ref.O_up.T.flatten()) + Assert.allclose(actual.O_down, spin_projectors.ref.O_down.T.flatten()) + Assert.allclose(actual.Fe_d_up, spin_projectors.ref.Fe_d_up.T.flatten()) + Assert.allclose(actual.Fe_d_down, spin_projectors.ref.Fe_d_down.T.flatten()) + + +def test_single_band_plot(single_band, Assert): + fig = single_band.plot(fermi_energy=single_band.ref.fermi_energy_argument) + assert fig.ylabel == "Energy (eV)" + assert len(fig.series) == 1 + assert fig.series[0].weight is None + Assert.allclose(fig.series[0].x, single_band.ref.kpoints.distances()) + Assert.allclose(fig.series[0].y, single_band.ref.bands.T) + + +def test_multiple_bands_plot(multiple_bands, Assert): + fig = multiple_bands.plot() + assert len(fig.series) == 1 # all bands in one plot + assert len(fig.series[0].x) == fig.series[0].y.shape[-1] + Assert.allclose(fig.series[0].y, multiple_bands.ref.bands.T) + + +def test_with_projectors_plot_default_width(with_projectors, Assert): + default_width = 0.5 + fig = with_projectors.plot(selection="Sr, p") + check_figure(fig, default_width, with_projectors.ref, Assert) + + +def test_with_projectors_plot_custom_width(with_projectors, Assert): + width = 0.1 + fig = with_projectors.plot(selection="Sr, p", width=width) + check_figure(fig, width, with_projectors.ref, Assert) + + +def test_spin_projectors_plot(spin_projectors, Assert): + reference = spin_projectors.ref + width = 0.05 + fig = spin_projectors.plot("O", width=width) + assert len(fig.series) == 2 + assert fig.series[0].label == "O_up" + check_data(fig.series[0], width, reference.bands_up, reference.O_up, Assert) + assert fig.series[1].label == "O_down" + check_data(fig.series[1], width, reference.bands_down, reference.O_down, Assert) + + +def check_figure(fig, weight, reference, Assert): + assert len(fig.series) == 2 + assert fig.series[0].label == "Sr" + assert fig.series[1].label == "p" + check_data(fig.series[0], weight, reference.bands, reference.Sr, Assert) + check_data(fig.series[1], weight, reference.bands, reference.p, Assert) + + +def check_data(series, weight, band, projection, Assert): + assert len(series.x) == series.y.shape[-1] + assert series.y.shape == series.weight.shape + Assert.allclose(series.y, band.T) + Assert.allclose(series.weight, weight * projection.T) + + +def test_spin_polarized_plot(spin_polarized, Assert): + fig = spin_polarized.plot() + assert len(fig.series) == 2 + assert fig.series[0].label == "up" + Assert.allclose(fig.series[0].y, spin_polarized.ref.bands_up.T) + assert fig.series[1].label == "down" + Assert.allclose(fig.series[1].y, spin_polarized.ref.bands_down.T) + + +def test_line_no_labels_plot(line_no_labels, Assert): + fig = line_no_labels.plot() + check_ticks(fig, line_no_labels.ref.kpoints, Assert) + reference_labels = ( + "$[0 0 0]$", + "$[0 0 \\frac{1}{2}]$", + "$[\\frac{1}{2} \\frac{1}{2} \\frac{1}{2}]$|$[0 0 0]$", + "$[\\frac{1}{2} \\frac{1}{2} 0]$", + "$[\\frac{1}{2} \\frac{1}{2} \\frac{1}{2}]$", + ) + assert tuple(fig.xticks.values()) == reference_labels + + +def test_line_with_labels_plot(line_with_labels, Assert): + fig = line_with_labels.plot() + check_ticks(fig, line_with_labels.ref.kpoints, Assert) + assert tuple(fig.xticks.values()) == (r"$\Gamma$", "", r"M|$\Gamma$", "Y", "M") + + +def test_noncollinear_plot(noncollinear_projectors, Assert): + default_width = 0.5 + fig = noncollinear_projectors.plot("total sigma_x") + assert len(fig.series) == 2 + assert fig.series[0].label == "total" + reference = noncollinear_projectors.ref + Assert.allclose(fig.series[0].weight, default_width * reference.total.T) + assert fig.series[0].marker is None + assert fig.series[1].label == "sigma_x" + Assert.allclose(fig.series[1].weight, reference.sigma_x.T) + assert fig.series[1].marker == "o" + assert fig.series[1].weight_mode == "color" + + +def check_ticks(fig, kpoints, Assert): + dists = kpoints.distances() + xticks = (*dists[:: kpoints.line_length()], dists[-1]) + Assert.allclose(list(fig.xticks.keys()), np.array(xticks)) + + +def test_plot_incorrect_width(with_projectors): + with pytest.raises(exception.IncorrectUsage): + with_projectors.plot("Sr", width="not a number") + + +@patch.object(Band, "to_graph") +def test_to_plotly(mock_plot, single_band): + fig = single_band.to_plotly("selection", width=0.2) + mock_plot.assert_called_once_with("selection", width=0.2) + graph = mock_plot.return_value + graph.to_plotly.assert_called_once() + assert fig == graph.to_plotly.return_value + + +def test_to_image(single_band): + check_to_image(single_band, None, "band.png") + custom_filename = "custom.jpg" + check_to_image(single_band, custom_filename, custom_filename) + + +def check_to_image(single_band, filename_argument, expected_filename): + with patch.object(Band, "to_plotly") as plot: + single_band.to_image("args", filename=filename_argument, key="word") + plot.assert_called_once_with("args", key="word") + fig = plot.return_value + fig.write_image.assert_called_once_with(single_band._path / expected_filename) + + +def test_band_selections(with_projectors): + actual = with_projectors.selections() + actual.pop("band") # remove band selections + assert actual == with_projectors.ref.selections + + +def test_to_quiver_with_incorrect_selection_raises_error(spin_texture_xy): + with pytest.raises(exception.IncorrectUsage): + spin_texture_xy.to_quiver("x") + with pytest.raises(exception.IncorrectUsage): + spin_texture_xy.to_quiver("x~y") + with pytest.raises(exception.IncorrectUsage): + spin_texture_xy.to_quiver("band=2") + + +@pytest.mark.parametrize( + "selection", + [ + "band=2(sigma_x~sigma_y)", + "Pb(sigma_1~sigma_2(band=1))", + "band=3(d(y~z))", + "p(5(band=1(z~x)))", + ], +) +def test_band_to_quiver(spin_texture, selection, Assert): + graph = spin_texture.to_quiver(selection) + assert graph.title == "Spin Texture" + assert len(graph) == 1 + series = graph.series[0] + Assert.allclose( + series.lattice.vectors, spin_texture.ref.expected_lattice, tolerance=1e6 + ) + assert series.label in spin_texture.ref.expected_data + Assert.allclose(series.data, spin_texture.ref.expected_data[series.label]) + + +@pytest.mark.parametrize("normal", [None, "x", "y", "z"]) +def test_band_to_quiver_normal(spin_texture_xy, normal, Assert): + graph = spin_texture_xy.to_quiver("band=1(x~z)", normal=normal) + assert len(graph) == 1 + quiver_plot = graph.series[0] + if normal == "x": + expected_lattice = [[-0.32989453, 1.48598944], [1.44773855, 0.47015754]] + elif normal == "y": + expected_lattice = [[1.44773855, 0.47015754], [-0.32989453, 1.48598944]] + elif normal == "z": + expected_lattice = [[1.52043113, 0.07269258], [0.07269258, 1.52043113]] + else: + expected_lattice = spin_texture_xy.ref.expected_lattice + Assert.allclose(quiver_plot.lattice.vectors, expected_lattice, tolerance=1e6) + + +@pytest.mark.parametrize( + "supercell, expected_supercell", [(None, (1, 1)), (2, (2, 2)), ([2, 4], (2, 4))] +) +def test_band_to_quiver_supercell( + spin_texture_xy, supercell, expected_supercell, Assert +): + graph = spin_texture_xy.to_quiver("band=3(sigma_2~sigma_1)", supercell=supercell) + assert len(graph) == 1 + quiver_plot = graph.series[0] + Assert.allclose(quiver_plot.supercell, expected_supercell) + + +def test_band_to_quiver_asymmetric_mesh(asymmetric_spin_texture, Assert): + """With a heavily asymmetric k-mesh (15x5x1), verify correct axis ordering.""" + graph = asymmetric_spin_texture.to_quiver("x~y(band=1)") + quiver_data = graph.series[0].data + assert quiver_data.shape[1] == 15 + assert quiver_data.shape[2] == 5 + Assert.allclose(quiver_data, asymmetric_spin_texture.ref.expected_data) + + +def test_band_to_quiver_rotation_projects_spin_vectors(rotated_spin_texture, Assert): + """Verify that spin vectors are rotated into the plot frame.""" + graph = rotated_spin_texture.to_quiver("x~y(band=1)") + quiver_data = graph.series[0].data + assert quiver_data.shape == (2, 4, 3) + # A pure Cartesian x-vector should project onto both plot axes for a rotated cell + assert not np.allclose(quiver_data[1], 0), ( + "Rotation is not applied: sigma_x should project onto both plot axes " + "for a 45-degree rotated cell" + ) + Assert.allclose(quiver_data, rotated_spin_texture.ref.expected_data) + + +def test_multiple_bands_print(multiple_bands, format_): + actual, _ = format_(multiple_bands) + reference = f""" +band data: + 48 k-points + 3 bands +no projectors + """.strip() + assert actual == {"text/plain": reference} + + +def test_line_no_labels_print(line_no_labels, format_): + actual, _ = format_(line_no_labels) + reference = f""" +band data: + 20 k-points + 3 bands +no projectors + """.strip() + assert actual == {"text/plain": reference} + + +def test_line_with_labels_print(line_with_labels, format_): + actual, _ = format_(line_with_labels) + reference = f""" +band data: + 20 k-points + 3 bands +no projectors + """.strip() + assert actual == {"text/plain": reference} + + +def test_spin_projectors_print(spin_projectors, format_): + actual, _ = format_(spin_projectors) + reference = f""" +spin polarized band data: + 48 k-points + 3 bands +{spin_projectors.ref.projectors_string} + """.strip() + assert actual == {"text/plain": reference} + + +def _check_to_database(_band): + handler = BandHandler.from_data(_band.ref.raw_data) + database_data: Band_DB = handler.to_database( + fermi_energy=getattr(_band.ref, "fermi_energy_argument", None) + ) + + assert isinstance(database_data, Band_DB) + + assert database_data.fermi_energy_raw == _band.ref.fermi_energy + assert database_data.fermi_energy == getattr( + _band.ref, "fermi_energy_argument", _band.ref.fermi_energy + ) + + if getattr(_band.ref, "num_occupied_bands", None) is not None: + assert database_data.num_occupied_bands == _band.ref.num_occupied_bands + elif ( + getattr(_band.ref, "occupations_up", None) is not None + and getattr(_band.ref, "occupations_down", None) is not None + ): + assert database_data.num_occupied_bands_up == _band.ref.num_occupied_bands_up + assert ( + database_data.num_occupied_bands_down == _band.ref.num_occupied_bands_down + ) + + for fld in fields(Band_DB): + if fld.name.startswith("num"): + assert getattr(database_data, fld.name) is None or isinstance( + getattr(database_data, fld.name), int + ), f"{fld.name} has unexpected type {type(getattr(database_data, fld.name))}: {getattr(database_data, fld.name)}" + else: + assert ( + getattr(database_data, fld.name) is None + or isinstance(getattr(database_data, fld.name), float) + or ( + fld.name.startswith("__") + and isinstance(getattr(database_data, fld.name), str) + ) + ), f"{fld.name} has unexpected type {type(getattr(database_data, fld.name))}: {getattr(database_data, fld.name)}" + + +def test_to_database_single_band(single_band): + _check_to_database(single_band) + + +def test_to_database_multiple_bands(multiple_bands): + _check_to_database(multiple_bands) + + +def test_to_database_spin_polarized(spin_polarized): + _check_to_database(spin_polarized) + + +def test_to_database_noncollinear_projectors(noncollinear_projectors): + _check_to_database(noncollinear_projectors) + + +def test_dispatcher_to_database_default(single_band): + """Dispatcher._to_database() must return {selection_name: handler_result}.""" + result = single_band._to_database() + assert isinstance(result, dict) + assert "band" in result + assert isinstance(result["band"], Band_DB) + + +def test_factory_methods(raw_data, check_factory_methods): + data = raw_data.band("multiple") + parameters = {"to_quiver": {"selection": "x~y(band=1)"}} + check_factory_methods(Band, data, parameters) diff --git a/tests/calculation/test_calculation_registry.py b/tests/calculation/test_calculation_registry.py index b173281da..f6bc5f4f3 100644 --- a/tests/calculation/test_calculation_registry.py +++ b/tests/calculation/test_calculation_registry.py @@ -1,197 +1,197 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -"""Tests for Calculation wiring to _REGISTRY and _source storage.""" - -import contextlib -import pathlib -from unittest.mock import patch - -import pytest - -from py4vasp import Calculation -from py4vasp._calculation.dispatch import ( - _REGISTRY, - DataSource, - DictSource, - FileSource, - Group, - quantity, -) - - -@contextlib.contextmanager -def _isolated_registry(): - """Restore _REGISTRY to its original state after the block.""" - saved = {k: dict(v) if isinstance(v, dict) else v for k, v in _REGISTRY.items()} - try: - yield - finally: - _REGISTRY.clear() - _REGISTRY.update(saved) - - -class _FakeDispatcher: - _quantity_name = "fake_qty" - - def __init__(self, source, quantity_name="fake_qty"): - self.source = source - self.quantity_name = quantity_name - - -class _FakeGroupDispatcher: - _quantity_name = "fake_member" - - def __init__(self, source, quantity_name="fake_member"): - self.source = source - self.quantity_name = quantity_name - - -class TestCalculationStoresSource: - def test_from_path_stores_file_source(self, tmp_path): - calc = Calculation.from_path(tmp_path) - assert isinstance(calc._source, FileSource) - assert calc._source.path == tmp_path.resolve() - - def test_from_path_file_source_has_no_file(self, tmp_path): - calc = Calculation.from_path(tmp_path) - assert calc._source._file is None - - def test_from_file_stores_file_source(self, tmp_path): - file = tmp_path / "vaspout.h5" - calc = Calculation.from_file(file) - assert isinstance(calc._source, FileSource) - - def test_from_file_source_path_is_parent_dir(self, tmp_path): - file = tmp_path / "vaspout.h5" - calc = Calculation.from_file(file) - assert calc._source.path == tmp_path.resolve() - - def test_from_file_source_file_is_forwarded(self, tmp_path): - file = tmp_path / "vaspout.h5" - calc = Calculation.from_file(file) - assert calc._source._file == file - - def test_path_property_still_works_after_from_path(self, tmp_path): - calc = Calculation.from_path(tmp_path) - assert calc.path() == tmp_path.resolve() - - def test_path_property_still_works_after_from_file(self, tmp_path): - file = tmp_path / "vaspout.h5" - calc = Calculation.from_file(file) - assert calc.path() == tmp_path.resolve() - - -class TestCalculationGetattr: - def test_getattr_top_level_returns_registered_dispatcher(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_qty") - class FakeDispatcher(_FakeDispatcher): - pass - - calc = Calculation.from_path(tmp_path) - result = calc.fake_qty - assert isinstance(result, FakeDispatcher) - - def test_getattr_dispatcher_receives_source(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_qty2") - class FakeDispatcher2(_FakeDispatcher): - _quantity_name = "fake_qty2" - - def __init__(self, source, quantity_name="fake_qty2"): - self.source = source - self.quantity_name = quantity_name - - calc = Calculation.from_path(tmp_path) - result = calc.fake_qty2 - assert result.source is calc._source - - def test_getattr_group_returns_group_instance(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_member", group="fake_group") - class FakeGroupDispatcher(_FakeGroupDispatcher): - pass - - calc = Calculation.from_path(tmp_path) - result = calc.fake_group - assert isinstance(result, Group) - - def test_getattr_group_member_instantiates_dispatcher(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_member2", group="fake_group2") - class FakeGroupDispatcher2(_FakeGroupDispatcher): - _quantity_name = "fake_member2" - - def __init__(self, source, quantity_name="fake_member2"): - self.source = source - self.quantity_name = quantity_name - - calc = Calculation.from_path(tmp_path) - result = calc.fake_group2.fake_member2 - assert isinstance(result, FakeGroupDispatcher2) - - def test_getattr_unknown_attribute_raises_attribute_error(self, tmp_path): - calc = Calculation.from_path(tmp_path) - with pytest.raises(AttributeError): - calc.this_quantity_does_not_exist_in_any_registry - - def test_getattr_dispatcher_source_path_matches_calculation_path(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_qty3") - class FakeDispatcher3(_FakeDispatcher): - _quantity_name = "fake_qty3" - - def __init__(self, source, quantity_name="fake_qty3"): - self.source = source - self.quantity_name = quantity_name - - calc = Calculation.from_path(tmp_path) - assert calc.fake_qty3.source.path == calc._path - - def test_new_arch_quantities_accessible_via_getattr(self): - """New-arch quantities registered via @quantity are accessible via __getattr__.""" - with patch("py4vasp.raw.access"): - calc = Calculation.from_path(".") - # 'energy' is a new-arch quantity wired via _REGISTRY/__getattr__ - assert hasattr(calc, "energy") - - -class TestCalculationDir: - def test_registry_quantities_appear_in_dir(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_dir_qty") - class FakeDirDispatcher(_FakeDispatcher): - _quantity_name = "fake_dir_qty" - - def __init__(self, source, quantity_name="fake_dir_qty"): - self.source = source - self.quantity_name = quantity_name - - calc = Calculation.from_path(tmp_path) - assert "fake_dir_qty" in dir(calc) - - def test_registry_groups_appear_in_dir(self, tmp_path): - with _isolated_registry(): - - @quantity("fake_dir_member", group="fake_dir_group") - class FakeDirGroupDispatcher(_FakeDispatcher): - _quantity_name = "fake_dir_member" - - def __init__(self, source, quantity_name="fake_dir_member"): - self.source = source - self.quantity_name = quantity_name - - calc = Calculation.from_path(tmp_path) - assert "fake_dir_group" in dir(calc) - - def test_existing_attributes_still_in_dir(self, tmp_path): - calc = Calculation.from_path(tmp_path) - names = dir(calc) - assert "from_path" in names - assert "from_file" in names +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +"""Tests for Calculation wiring to _REGISTRY and _source storage.""" + +import contextlib +import pathlib +from unittest.mock import patch + +import pytest + +from py4vasp import Calculation +from py4vasp._calculation.dispatch import ( + _REGISTRY, + DataSource, + DictSource, + FileSource, + Group, + quantity, +) + + +@contextlib.contextmanager +def _isolated_registry(): + """Restore _REGISTRY to its original state after the block.""" + saved = {k: dict(v) if isinstance(v, dict) else v for k, v in _REGISTRY.items()} + try: + yield + finally: + _REGISTRY.clear() + _REGISTRY.update(saved) + + +class _FakeDispatcher: + _quantity_name = "fake_qty" + + def __init__(self, source, quantity_name="fake_qty"): + self.source = source + self.quantity_name = quantity_name + + +class _FakeGroupDispatcher: + _quantity_name = "fake_member" + + def __init__(self, source, quantity_name="fake_member"): + self.source = source + self.quantity_name = quantity_name + + +class TestCalculationStoresSource: + def test_from_path_stores_file_source(self, tmp_path): + calc = Calculation.from_path(tmp_path) + assert isinstance(calc._source, FileSource) + assert calc._source.path == tmp_path.resolve() + + def test_from_path_file_source_has_no_file(self, tmp_path): + calc = Calculation.from_path(tmp_path) + assert calc._source._file is None + + def test_from_file_stores_file_source(self, tmp_path): + file = tmp_path / "vaspout.h5" + calc = Calculation.from_file(file) + assert isinstance(calc._source, FileSource) + + def test_from_file_source_path_is_parent_dir(self, tmp_path): + file = tmp_path / "vaspout.h5" + calc = Calculation.from_file(file) + assert calc._source.path == tmp_path.resolve() + + def test_from_file_source_file_is_forwarded(self, tmp_path): + file = tmp_path / "vaspout.h5" + calc = Calculation.from_file(file) + assert calc._source._file == file + + def test_path_property_still_works_after_from_path(self, tmp_path): + calc = Calculation.from_path(tmp_path) + assert calc.path() == tmp_path.resolve() + + def test_path_property_still_works_after_from_file(self, tmp_path): + file = tmp_path / "vaspout.h5" + calc = Calculation.from_file(file) + assert calc.path() == tmp_path.resolve() + + +class TestCalculationGetattr: + def test_getattr_top_level_returns_registered_dispatcher(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_qty") + class FakeDispatcher(_FakeDispatcher): + pass + + calc = Calculation.from_path(tmp_path) + result = calc.fake_qty + assert isinstance(result, FakeDispatcher) + + def test_getattr_dispatcher_receives_source(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_qty2") + class FakeDispatcher2(_FakeDispatcher): + _quantity_name = "fake_qty2" + + def __init__(self, source, quantity_name="fake_qty2"): + self.source = source + self.quantity_name = quantity_name + + calc = Calculation.from_path(tmp_path) + result = calc.fake_qty2 + assert result.source is calc._source + + def test_getattr_group_returns_group_instance(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_member", group="fake_group") + class FakeGroupDispatcher(_FakeGroupDispatcher): + pass + + calc = Calculation.from_path(tmp_path) + result = calc.fake_group + assert isinstance(result, Group) + + def test_getattr_group_member_instantiates_dispatcher(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_member2", group="fake_group2") + class FakeGroupDispatcher2(_FakeGroupDispatcher): + _quantity_name = "fake_member2" + + def __init__(self, source, quantity_name="fake_member2"): + self.source = source + self.quantity_name = quantity_name + + calc = Calculation.from_path(tmp_path) + result = calc.fake_group2.fake_member2 + assert isinstance(result, FakeGroupDispatcher2) + + def test_getattr_unknown_attribute_raises_attribute_error(self, tmp_path): + calc = Calculation.from_path(tmp_path) + with pytest.raises(AttributeError): + calc.this_quantity_does_not_exist_in_any_registry + + def test_getattr_dispatcher_source_path_matches_calculation_path(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_qty3") + class FakeDispatcher3(_FakeDispatcher): + _quantity_name = "fake_qty3" + + def __init__(self, source, quantity_name="fake_qty3"): + self.source = source + self.quantity_name = quantity_name + + calc = Calculation.from_path(tmp_path) + assert calc.fake_qty3.source.path == calc._path + + def test_new_arch_quantities_accessible_via_getattr(self): + """New-arch quantities registered via @quantity are accessible via __getattr__.""" + with patch("py4vasp.raw.access"): + calc = Calculation.from_path(".") + # 'energy' is a new-arch quantity wired via _REGISTRY/__getattr__ + assert hasattr(calc, "energy") + + +class TestCalculationDir: + def test_registry_quantities_appear_in_dir(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_dir_qty") + class FakeDirDispatcher(_FakeDispatcher): + _quantity_name = "fake_dir_qty" + + def __init__(self, source, quantity_name="fake_dir_qty"): + self.source = source + self.quantity_name = quantity_name + + calc = Calculation.from_path(tmp_path) + assert "fake_dir_qty" in dir(calc) + + def test_registry_groups_appear_in_dir(self, tmp_path): + with _isolated_registry(): + + @quantity("fake_dir_member", group="fake_dir_group") + class FakeDirGroupDispatcher(_FakeDispatcher): + _quantity_name = "fake_dir_member" + + def __init__(self, source, quantity_name="fake_dir_member"): + self.source = source + self.quantity_name = quantity_name + + calc = Calculation.from_path(tmp_path) + assert "fake_dir_group" in dir(calc) + + def test_existing_attributes_still_in_dir(self, tmp_path): + calc = Calculation.from_path(tmp_path) + names = dir(calc) + assert "from_path" in names + assert "from_file" in names diff --git a/tests/calculation/test_contcar.py b/tests/calculation/test_contcar.py index 033f0d6a9..f80ca4ce6 100644 --- a/tests/calculation/test_contcar.py +++ b/tests/calculation/test_contcar.py @@ -1,131 +1,131 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import types - -import pytest - -from py4vasp._calculation._CONTCAR import CONTCAR as _CONTCAR -from py4vasp._calculation._CONTCAR import CONTCARHandler -from py4vasp._calculation.structure import Structure -from py4vasp._raw.data_db import CONTCAR_DB - -REF_Sr2TiO4 = """\ -Sr2TiO4 -6.9229000000000003 - 1.0000000000000000 0.0000000000000000 0.0000000000000000 - 0.6781122097386930 0.7349583872510080 0.0000000000000000 - -0.8390553410420490 -0.3674788590908430 0.4011800378743010 -Sr Ti O -2 1 4 -Direct - 0.6452900000000000 0.6452900000000000 0.0000000000000000 - 0.3547100000000000 0.3547100000000000 0.0000000000000000 - 0.0000000000000000 0.0000000000000000 0.0000000000000000 - 0.8417800000000000 0.8417800000000000 0.0000000000000000 - 0.1582300000000000 0.1582300000000000 0.0000000000000000 - 0.5000000000000000 0.0000000000000000 0.5000000000000000 - 0.0000000000000000 0.5000000000000000 0.5000000000000000""" - -REF_Fe3O4 = """\ -Fe3O4 -1.0000000000000000 - 5.1941269999999999 0.0000000000000000 0.0000000000000000 - 0.0000000000000000 3.0893880000000000 0.0000000000000000 - -1.3770129362480001 0.0000000000000000 5.0950563617919995 -Fe O -3 4 -Selective dynamics -Direct - 0.0100000000000000 0.0100000000000000 0.0100000000000000 T F T - 0.5100000000000000 0.0100000000000000 0.5100000000000000 F T F - 0.0100000000000000 0.5100000000000000 0.5100000000000000 T F T - 0.7974500000000000 0.0100000000000000 0.2915200000000000 F T F - 0.2731000000000000 0.5100000000000000 0.2861100000000000 T F T - 0.2225500000000000 0.0100000000000000 0.7284800000000000 F T F - 0.7469000000000000 0.5100000000000000 0.7338900000000000 T F T -Lattice velocities and vectors -1 - 2.39789553e+00 -3.00000000e-01 -3.00000000e-01 - -3.00000000e-01 6.54431821e-01 -3.00000000e-01 - -1.10383537e-01 -3.00000000e-01 2.29595993e+00 - 5.19412700e+00 0.00000000e+00 0.00000000e+00 - 0.00000000e+00 3.08938800e+00 0.00000000e+00 - -1.37701294e+00 0.00000000e+00 5.09505636e+00 -Cartesian - 0.00000000e+00 1.00000000e+00 1.41421356e+00 - 1.73205081e+00 2.00000000e+00 2.23606798e+00 - 2.44948974e+00 2.64575131e+00 2.82842712e+00 - 3.00000000e+00 3.16227766e+00 3.31662479e+00 - 3.46410162e+00 3.60555128e+00 3.74165739e+00 - 3.87298335e+00 4.00000000e+00 4.12310563e+00 - 4.24264069e+00 4.35889894e+00 4.47213595e+00""" - - -@pytest.fixture(params=["Sr2TiO4", "Fe3O4"]) -def CONTCAR(raw_data, request): - selection = request.param - raw_contcar = raw_data.CONTCAR(selection) - contcar = _CONTCAR.from_data(raw_contcar) - contcar.ref = types.SimpleNamespace() - structure = Structure.from_data(raw_data.structure(selection))[-1] - contcar.ref.structure = structure - contcar.ref.system = selection - contcar.ref.selective_dynamics = raw_contcar.selective_dynamics - contcar.ref.lattice_velocities = raw_contcar.lattice_velocities - contcar.ref.ion_velocities = raw_contcar.ion_velocities - contcar.ref.string = REF_Sr2TiO4 if selection == "Sr2TiO4" else REF_Fe3O4 - contcar.ref.raw_data = raw_contcar - return contcar - - -class OptionalOutputCheck: - def __init__(self, dict_, Assert): - self.dict_ = dict_ - self.Assert = Assert - - def element_agrees(self, key, reference): - if reference.is_none(): - assert key not in self.dict_ - else: - self.Assert.allclose(self.dict_[key], reference) - - -def test_read(CONTCAR, Assert): - actual = CONTCAR.read() - expected = CONTCAR.ref.structure.read() - Assert.allclose(actual["lattice_vectors"], expected["lattice_vectors"]) - Assert.allclose(actual["positions"], expected["positions"]) - assert actual["elements"] == expected["elements"] - assert actual["names"] == expected["names"] - assert actual["system"] == CONTCAR.ref.system - check = OptionalOutputCheck(actual, Assert) - check.element_agrees("selective_dynamics", CONTCAR.ref.selective_dynamics) - check.element_agrees("lattice_velocities", CONTCAR.ref.lattice_velocities) - check.element_agrees("ion_velocities", CONTCAR.ref.ion_velocities) - - -@pytest.mark.parametrize("supercell", [None, 2, (3, 2, 1)]) -def test_plot(CONTCAR, supercell, Assert): - structure_view = CONTCAR.ref.structure.plot(supercell) - view = CONTCAR.plot(supercell) if supercell else CONTCAR.plot() - Assert.same_structure_view(view, structure_view) - view = CONTCAR.to_view(supercell) if supercell else CONTCAR.to_view() - Assert.same_structure_view(view, structure_view) - - -def test_print(CONTCAR, format_): - actual, _ = format_(CONTCAR) - assert actual == {"text/plain": CONTCAR.ref.string} - - -def test_factory_methods(raw_data, check_factory_methods): - raw_contcar = raw_data.CONTCAR("Sr2TiO4") - check_factory_methods(_CONTCAR, raw_contcar) - - -def test_to_database(CONTCAR): - handler = CONTCARHandler.from_data(CONTCAR.ref.raw_data) - db_dict: CONTCAR_DB = handler.to_database() - - assert db_dict.system == CONTCAR.ref.system - assert isinstance(db_dict.system, (str, type(None))) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import types + +import pytest + +from py4vasp._calculation._CONTCAR import CONTCAR as _CONTCAR +from py4vasp._calculation._CONTCAR import CONTCARHandler +from py4vasp._calculation.structure import Structure +from py4vasp._raw.data_db import CONTCAR_DB + +REF_Sr2TiO4 = """\ +Sr2TiO4 +6.9229000000000003 + 1.0000000000000000 0.0000000000000000 0.0000000000000000 + 0.6781122097386930 0.7349583872510080 0.0000000000000000 + -0.8390553410420490 -0.3674788590908430 0.4011800378743010 +Sr Ti O +2 1 4 +Direct + 0.6452900000000000 0.6452900000000000 0.0000000000000000 + 0.3547100000000000 0.3547100000000000 0.0000000000000000 + 0.0000000000000000 0.0000000000000000 0.0000000000000000 + 0.8417800000000000 0.8417800000000000 0.0000000000000000 + 0.1582300000000000 0.1582300000000000 0.0000000000000000 + 0.5000000000000000 0.0000000000000000 0.5000000000000000 + 0.0000000000000000 0.5000000000000000 0.5000000000000000""" + +REF_Fe3O4 = """\ +Fe3O4 +1.0000000000000000 + 5.1941269999999999 0.0000000000000000 0.0000000000000000 + 0.0000000000000000 3.0893880000000000 0.0000000000000000 + -1.3770129362480001 0.0000000000000000 5.0950563617919995 +Fe O +3 4 +Selective dynamics +Direct + 0.0100000000000000 0.0100000000000000 0.0100000000000000 T F T + 0.5100000000000000 0.0100000000000000 0.5100000000000000 F T F + 0.0100000000000000 0.5100000000000000 0.5100000000000000 T F T + 0.7974500000000000 0.0100000000000000 0.2915200000000000 F T F + 0.2731000000000000 0.5100000000000000 0.2861100000000000 T F T + 0.2225500000000000 0.0100000000000000 0.7284800000000000 F T F + 0.7469000000000000 0.5100000000000000 0.7338900000000000 T F T +Lattice velocities and vectors +1 + 2.39789553e+00 -3.00000000e-01 -3.00000000e-01 + -3.00000000e-01 6.54431821e-01 -3.00000000e-01 + -1.10383537e-01 -3.00000000e-01 2.29595993e+00 + 5.19412700e+00 0.00000000e+00 0.00000000e+00 + 0.00000000e+00 3.08938800e+00 0.00000000e+00 + -1.37701294e+00 0.00000000e+00 5.09505636e+00 +Cartesian + 0.00000000e+00 1.00000000e+00 1.41421356e+00 + 1.73205081e+00 2.00000000e+00 2.23606798e+00 + 2.44948974e+00 2.64575131e+00 2.82842712e+00 + 3.00000000e+00 3.16227766e+00 3.31662479e+00 + 3.46410162e+00 3.60555128e+00 3.74165739e+00 + 3.87298335e+00 4.00000000e+00 4.12310563e+00 + 4.24264069e+00 4.35889894e+00 4.47213595e+00""" + + +@pytest.fixture(params=["Sr2TiO4", "Fe3O4"]) +def CONTCAR(raw_data, request): + selection = request.param + raw_contcar = raw_data.CONTCAR(selection) + contcar = _CONTCAR.from_data(raw_contcar) + contcar.ref = types.SimpleNamespace() + structure = Structure.from_data(raw_data.structure(selection))[-1] + contcar.ref.structure = structure + contcar.ref.system = selection + contcar.ref.selective_dynamics = raw_contcar.selective_dynamics + contcar.ref.lattice_velocities = raw_contcar.lattice_velocities + contcar.ref.ion_velocities = raw_contcar.ion_velocities + contcar.ref.string = REF_Sr2TiO4 if selection == "Sr2TiO4" else REF_Fe3O4 + contcar.ref.raw_data = raw_contcar + return contcar + + +class OptionalOutputCheck: + def __init__(self, dict_, Assert): + self.dict_ = dict_ + self.Assert = Assert + + def element_agrees(self, key, reference): + if reference.is_none(): + assert key not in self.dict_ + else: + self.Assert.allclose(self.dict_[key], reference) + + +def test_read(CONTCAR, Assert): + actual = CONTCAR.read() + expected = CONTCAR.ref.structure.read() + Assert.allclose(actual["lattice_vectors"], expected["lattice_vectors"]) + Assert.allclose(actual["positions"], expected["positions"]) + assert actual["elements"] == expected["elements"] + assert actual["names"] == expected["names"] + assert actual["system"] == CONTCAR.ref.system + check = OptionalOutputCheck(actual, Assert) + check.element_agrees("selective_dynamics", CONTCAR.ref.selective_dynamics) + check.element_agrees("lattice_velocities", CONTCAR.ref.lattice_velocities) + check.element_agrees("ion_velocities", CONTCAR.ref.ion_velocities) + + +@pytest.mark.parametrize("supercell", [None, 2, (3, 2, 1)]) +def test_plot(CONTCAR, supercell, Assert): + structure_view = CONTCAR.ref.structure.plot(supercell) + view = CONTCAR.plot(supercell) if supercell else CONTCAR.plot() + Assert.same_structure_view(view, structure_view) + view = CONTCAR.to_view(supercell) if supercell else CONTCAR.to_view() + Assert.same_structure_view(view, structure_view) + + +def test_print(CONTCAR, format_): + actual, _ = format_(CONTCAR) + assert actual == {"text/plain": CONTCAR.ref.string} + + +def test_factory_methods(raw_data, check_factory_methods): + raw_contcar = raw_data.CONTCAR("Sr2TiO4") + check_factory_methods(_CONTCAR, raw_contcar) + + +def test_to_database(CONTCAR): + handler = CONTCARHandler.from_data(CONTCAR.ref.raw_data) + db_dict: CONTCAR_DB = handler.to_database() + + assert db_dict.system == CONTCAR.ref.system + assert isinstance(db_dict.system, (str, type(None))) diff --git a/tests/calculation/test_default_calculation.py b/tests/calculation/test_default_calculation.py index e394d6c05..a44185a15 100644 --- a/tests/calculation/test_default_calculation.py +++ b/tests/calculation/test_default_calculation.py @@ -1,8 +1,12 @@ # Copyright © VASP Software GmbH, # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import contextlib + import pytest -from py4vasp import Calculation, calculation +from py4vasp import Calculation, calculation, demo +from py4vasp._raw.schema import DEFAULT_SELECTION +from py4vasp._util import loadable def test_access_of_attributes(): @@ -27,3 +31,256 @@ def test_assigning_to_input_file(tmp_path, monkeypatch): with open("INCAR", "r") as file: actual = file.read() assert actual == expected + + +def test_selections_on_empty_path(tmp_path): + # Default (only_available=False) still returns all schema-defined quantities on empty path + calc = Calculation.from_path(tmp_path) + full = calc.selections() + assert "band" in full + assert "bandgap" in full + assert full["band"] == ["default", "kpoints_opt", "kpoints_wan"] + + +def test_selections_on_demo_calculation(tmp_path): + calc = demo.calculation(tmp_path / "demo_calculation") + actual = calc.selections() + # Default now returns all public quantities with schema-defined selections + assert "band" in actual + assert "bandgap" in actual # included even without data + assert "density" in actual + assert "structure" in actual + # the result is sorted by quantity name + assert list(actual) == sorted(actual) + + +def test_selections_loadable_excludes_selections_that_do_not_load(tmp_path): + # Use only_available=True to get only loadable quantities + calc = demo.calculation(tmp_path / "demo_calculation") + actual = calc.selections(only_available=True) + # current_density cannot be read without specifying a cut plane -> excluded + assert "current_density" not in actual + # a non-default source of a not-yet-migrated quantity cannot be addressed + assert actual["structure"] == ["default"] + + +def test_selections_evaluable(tmp_path): + calculation = demo.calculation(tmp_path / "demo_calculation") + # selections with method parameter should work and return available sources + viewable = calculation.selections(method="to_view") + assert isinstance(viewable, dict) + for quantity, sources in viewable.items(): + assert isinstance(sources, list) + assert all(isinstance(s, str) for s in sources) + + +def test_selections_includes_quantities_without_data(tmp_path): + # Default now includes all quantities; quantities without data have empty selections + calc = demo.calculation(tmp_path / "demo_calculation") + actual = calc.selections() + included = ( + "bandgap", + "born_effective_charge", + "dielectric_function", + "dielectric_tensor", + "elastic_modulus", + "internal_strain", + "piezoelectric_tensor", + "polarization", + ) + for quantity in included: + assert quantity in actual + + +def test_selections_filtered_by_method(tmp_path): + calc = demo.calculation(tmp_path / "demo_calculation") + viewable = calc.selections(method="to_view") + full = calc.selections() + # only quantities implementing the method are reported + assert set(viewable) <= set(full) + assert viewable.keys() >= {"density", "potential", "structure"} + # quantities without a to_view method are excluded + for quantity in ("band", "dos", "energy", "stress"): + assert quantity not in viewable + + +def test_selections_with_method_on_empty_path(tmp_path): + # Default (only_available=False) with method filter still returns quantities implementing the method + calc = Calculation.from_path(tmp_path) + result = calc.selections(method="to_view") + assert "density" in result + assert "structure" in result + assert "band" not in result + + +def test_selections_with_only_available_true(tmp_path): + calc = demo.calculation(tmp_path / "demo_calculation") + loadable = calc.selections(only_available=True) + full = calc.selections(only_available=False) + # loadable quantities should be a subset of all quantities + assert set(loadable) <= set(full) + # quantities without loadable data should not appear in loadable result + absent_in_loadable = { + "bandgap", + "born_effective_charge", + "dielectric_function", + "dielectric_tensor", + "elastic_modulus", + "internal_strain", + "piezoelectric_tensor", + "polarization", + } + for quantity in absent_in_loadable: + assert quantity not in loadable + assert quantity in full + assert "default" in full[quantity] + assert full[quantity] + + +def test_selections_on_empty_path_returns_all(tmp_path): + # Default (only_available=False) returns schema-defined selections even without data + calc = Calculation.from_path(tmp_path) + full = calc.selections() + + assert full["band"] == ["default", "kpoints_opt", "kpoints_wan"] + assert "default" in full["structure"] + assert "final" in full["structure"] + assert "poscar" in full["structure"] + assert full["exciton.density"] == ["default"] + + +def test_selections_on_empty_path_only_available_true(tmp_path): + # With only_available=True on empty path, nothing loads + calc = Calculation.from_path(tmp_path) + assert calc.selections(only_available=True) == {} + + +def test_selections_with_method_filters_by_implementation(tmp_path): + calc = demo.calculation(tmp_path / "demo_calculation") + # Default: all quantities implementing to_view with schema selections + full_view = calc.selections(method="to_view") + + assert "band" not in full_view + assert "dos" not in full_view + assert full_view["density"] == ["default", "tau"] + assert "default" in full_view["structure"] + + +def test_all_quantities_implement_read(tmp_path): + """Verify all quantities implement the read() method.""" + calc = demo.calculation(tmp_path / "demo_calculation") + all_quantities = calc.selections(only_available=False) + for quantity_name in all_quantities: + assert loadable.implements_method( + calc, quantity_name, "read" + ), f"{quantity_name} does not implement read()" + + +def test_selections_only_available_false_does_not_load_data(tmp_path, monkeypatch): + # only_available=False must never attempt to load data + calc = demo.calculation(tmp_path / "demo_calculation") + + def _fail(*_, **__): + raise AssertionError("loadable_sources should not be called") + + monkeypatch.setattr(loadable, "loadable_sources", _fail) + + # neither the default nor the method-filtered call should load any data + assert calc.selections() + assert calc.selections(method="to_view") + + +def test_selections_collects_out_of_memory_errors(tmp_path, monkeypatch, capsys): + # A quantity whose data is available but too large to load raises MemoryError. + # selections() should not crash, should omit that quantity, and report the error. + calc = demo.calculation(tmp_path / "demo_calculation") + + monkeypatch.setattr(loadable, "_schema_satisfied", lambda *_, **__: None) + real_invoke = loadable._invoke + + def _invoke_with_oom( + calculation, + call_name, + method_name, + source_name, + legacy_quantities, + convention=None, + ): + if call_name == "density": + raise MemoryError("Unable to allocate 5.00 GiB for array") + return real_invoke( + calculation, + call_name, + method_name, + source_name, + legacy_quantities, + convention, + ) + + monkeypatch.setattr(loadable, "_invoke", _invoke_with_oom) + + result = calc.selections(only_available=True) + + # the oversized quantity is excluded from the loadable result + assert "density" not in result + # other quantities are still inspected and the call returns normally + assert isinstance(result, dict) + # the out-of-memory error is reported once at the end without crashing + captured = capsys.readouterr() + assert "density" in captured.out + assert "out of memory" in captured.out.lower() + assert "5.00 GiB" in captured.out + + +def test_selections_default_does_not_report_out_of_memory( + tmp_path, monkeypatch, capsys +): + # only_available=False never loads data, so it cannot trigger an OOM report + calc = demo.calculation(tmp_path / "demo_calculation") + + def _fail(*_, **__): + raise AssertionError("loadable_sources should not be called") + + monkeypatch.setattr(loadable, "loadable_sources", _fail) + + calc.selections() + captured = capsys.readouterr() + assert "out of memory" not in captured.out.lower() + + +def test_confirm_read_uses_public_call_name_for_fallback(tmp_path, monkeypatch): + calc = demo.calculation(tmp_path / "demo_calculation") + captured = {} + + monkeypatch.setattr(loadable, "_schema_satisfied", lambda *_, **__: None) + + def _record_invoke( + calculation, + call_name, + method_name, + source_name, + legacy_quantities, + convention=None, + ): + captured["call_name"] = call_name + captured["method_name"] = method_name + captured["source_name"] = source_name + return True + + monkeypatch.setattr(loadable, "_invoke", _record_invoke) + + with contextlib.ExitStack() as stack: + assert loadable._confirm_read( + calc, + "exciton.density", + "exciton_density", + DEFAULT_SELECTION, + open_files={}, + stack=stack, + cache={}, + legacy_quantities=set(), + ) + + assert captured["call_name"] == "exciton.density" + assert captured["method_name"] == "read" + assert captured["source_name"] == DEFAULT_SELECTION diff --git a/tests/calculation/test_partial_density.py b/tests/calculation/test_partial_density.py index 7b28b86bb..d7ae2950e 100644 --- a/tests/calculation/test_partial_density.py +++ b/tests/calculation/test_partial_density.py @@ -1,414 +1,414 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -import dataclasses -import types - -import numpy as np -import pytest - -from py4vasp import _config -from py4vasp._calculation.partial_density import PartialDensity -from py4vasp._calculation.structure import Structure, StructureHandler -from py4vasp._util.slicing import plane -from py4vasp.exception import IncorrectUsage, NoData, NotImplemented - - -@pytest.fixture( - params=[ - "no splitting no spin", - "no splitting no spin Ca3AsBr3", - "no splitting no spin Sr2TiO4", - "no splitting no spin CaAs3_110", - "split_bands", - "split_bands and spin_polarized", - "split_bands and spin_polarized Ca3AsBr3", - "split_bands and spin_polarized Sr2TiO4", - "split_kpoints", - "split_kpoints and spin_polarized", - "split_kpoints and spin_polarized Ca3AsBr3", - "split_kpoints and spin_polarized Sr2TiO4", - "spin_polarized", - "spin_polarized Ca3AsBr3", - "spin_polarized Sr2TiO4", - "split_bands and split_kpoints", - "split_bands and split_kpoints and spin_polarized", - "split_bands and split_kpoints and spin_polarized Ca3AsBr3", - "split_bands and split_kpoints and spin_polarized Sr2TiO4", - ] -) -def AnyPartialDensity(raw_data, request): - return make_reference_partial_density(raw_data, request.param) - - -@pytest.fixture -def NonSplitPartialDensity(raw_data): - return make_reference_partial_density(raw_data, "no splitting no spin") - - -@pytest.fixture -def PolarizedNonSplitPartialDensity(raw_data): - return make_reference_partial_density(raw_data, "spin_polarized") - - -@pytest.fixture -def PolarizedNonSplitPartialDensityCa3AsBr3(raw_data): - return make_reference_partial_density(raw_data, "spin_polarized Ca3AsBr3") - - -@pytest.fixture -def NonSplitPartialDensityCaAs3_110(raw_data): - return make_reference_partial_density(raw_data, "CaAs3_110") - - -@pytest.fixture -def NonSplitPartialDensityNi_100(raw_data): - return make_reference_partial_density(raw_data, "Ni100") - - -@pytest.fixture -def PolarizedNonSplitPartialDensitySr2TiO4(raw_data): - return make_reference_partial_density(raw_data, "spin_polarized Sr2TiO4") - - -@pytest.fixture -def NonPolarizedBandSplitPartialDensity(raw_data): - return make_reference_partial_density(raw_data, "split_bands") - - -@pytest.fixture -def PolarizedAllSplitPartialDensity(raw_data): - return make_reference_partial_density( - raw_data, "split_bands and split_kpoints and spin_polarized" - ) - - -@pytest.fixture(params=["up", "down", "total"]) -def spin(request): - return request.param - - -def make_reference_partial_density(raw_data, selection): - raw_partial_density = raw_data.partial_density(selection=selection) - parchg = PartialDensity.from_data(raw_partial_density) - parchg.ref = types.SimpleNamespace() - parchg.ref.structure = Structure.from_data(raw_partial_density.structure) - parchg.ref.structure_handler = StructureHandler.from_data( - raw_partial_density.structure - ) - parchg.ref.plane_vectors = plane( - cell=parchg.ref.structure.lattice_vectors(), - cut="c", - normal="z", - ) - parchg.ref.partial_density = raw_partial_density.partial_charge - parchg.ref.bands = raw_partial_density.bands - parchg.ref.kpoints = raw_partial_density.kpoints - parchg.ref.grid = raw_partial_density.grid - return parchg - - -def test_read(AnyPartialDensity, Assert, not_core): - actual = AnyPartialDensity.read() - expected = AnyPartialDensity.ref - Assert.allclose(actual["bands"], expected.bands) - Assert.allclose(actual["kpoints"], expected.kpoints) - Assert.allclose(actual["grid"], expected.grid) - expected_density = np.squeeze(np.asarray(expected.partial_density).T) - Assert.allclose(actual["partial_density"], expected_density) - Assert.same_structure(actual["structure"], expected.structure.read()) - - -def test_stoichiometry(AnyPartialDensity, not_core): - actual = AnyPartialDensity._stoichiometry() - expected = str(AnyPartialDensity.ref.structure_handler._stoichiometry()) - assert actual == expected - - -def test_bands(AnyPartialDensity, Assert, not_core): - actual = AnyPartialDensity.bands() - expected = AnyPartialDensity.ref.bands - Assert.allclose(actual, expected) - - -def test_kpoints(AnyPartialDensity, Assert, not_core): - actual = AnyPartialDensity.kpoints() - expected = AnyPartialDensity.ref.kpoints - Assert.allclose(actual, expected) - - -def test_grid(AnyPartialDensity, Assert, not_core): - actual = AnyPartialDensity.grid() - expected = AnyPartialDensity.ref.grid - Assert.allclose(actual, expected) - - -@pytest.mark.parametrize("selection", (None, "total", "up", "down")) -def test_plot(PolarizedNonSplitPartialDensity, selection, Assert, not_core): - if selection is not None: - view = PolarizedNonSplitPartialDensity.plot(selection) - expected_density = PolarizedNonSplitPartialDensity.to_numpy(selection) - expected_label = selection - else: - view = PolarizedNonSplitPartialDensity.plot() - expected_density = PolarizedNonSplitPartialDensity.to_numpy() - expected_label = "total" - reference = PolarizedNonSplitPartialDensity.ref - Assert.same_structure_view(view, reference.structure.plot()) - assert len(view.grid_scalars) == 1 - grid_scalar = view.grid_scalars[0] - assert grid_scalar.label == expected_label - assert grid_scalar.quantity.ndim == 4 - Assert.allclose(grid_scalar.quantity, expected_density) - assert len(grid_scalar.isosurfaces) == 1 - isosurface = grid_scalar.isosurfaces[0] - assert isosurface.isolevel == 0.2 - assert isosurface.color == _config.VASP_COLORS["cyan"] - assert isosurface.opacity == 0.6 - - -def test_plot_with_options(NonSplitPartialDensity, Assert, not_core): - view = NonSplitPartialDensity.plot(supercell=(1, 2, 3), isolevel=0.6, color="red") - Assert.allclose(view.supercell, (1, 2, 3)) - assert len(view.grid_scalars) == 1 - grid_scalar = view.grid_scalars[0] - assert len(grid_scalar.isosurfaces) == 1 - isosurface = grid_scalar.isosurfaces[0] - assert isosurface.isolevel == 0.6 - assert isosurface.color == "red" - - -def test_non_split_to_numpy(PolarizedNonSplitPartialDensity, Assert, not_core): - actual = PolarizedNonSplitPartialDensity.to_numpy("total") - expected = PolarizedNonSplitPartialDensity.ref.partial_density - Assert.allclose(actual, expected[0, 0, 0].T) - - actual = PolarizedNonSplitPartialDensity.to_numpy("up") - Assert.allclose(actual, 0.5 * (expected[0, 0, 0].T + expected[0, 0, 1].T)) - - actual = PolarizedNonSplitPartialDensity.to_numpy("down") - Assert.allclose(actual, 0.5 * (expected[0, 0, 0].T - expected[0, 0, 1].T)) - - -def test_split_to_numpy(PolarizedAllSplitPartialDensity, Assert, not_core): - bands = PolarizedAllSplitPartialDensity.ref.bands - kpoints = PolarizedAllSplitPartialDensity.ref.kpoints - for band_index, band in enumerate(bands): - for kpoint_index, kpoint in enumerate(kpoints): - actual = PolarizedAllSplitPartialDensity.to_numpy( - band=band, kpoint=kpoint, selection="total" - ) - expected = PolarizedAllSplitPartialDensity.ref.partial_density - Assert.allclose(actual, np.asarray(expected)[kpoint_index, band_index, 0].T) - - msg = f"Band {max(bands) + 1} not found in the bands array." - with pytest.raises(NoData) as excinfo: - PolarizedAllSplitPartialDensity.to_numpy( - band=max(bands) + 1, kpoint=max(kpoints), selection="up" - ) - assert msg in str(excinfo.value) - - msg = f"K-point {min(kpoints) - 1} not found in the kpoints array." - with pytest.raises(NoData) as excinfo: - PolarizedAllSplitPartialDensity.to_numpy( - band=min(bands), kpoint=min(kpoints) - 1, selection="down" - ) - assert msg in str(excinfo.value) - - -def test_non_polarized_to_numpy(NonSplitPartialDensity, spin, Assert, not_core): - actual = NonSplitPartialDensity.to_numpy(selection=spin) - expected = NonSplitPartialDensity.ref.partial_density - Assert.allclose(actual, np.asarray(expected).T[:, :, :, 0, 0, 0]) - - -def test_split_bands_to_numpy( - NonPolarizedBandSplitPartialDensity, spin, Assert, not_core -): - bands = NonPolarizedBandSplitPartialDensity.ref.bands - for band_index, band in enumerate(bands): - actual = NonPolarizedBandSplitPartialDensity.to_numpy(spin, band=band) - expected = NonPolarizedBandSplitPartialDensity.ref.partial_density - Assert.allclose(actual, np.asarray(expected).T[:, :, :, 0, band_index, 0]) - - -def test_to_stm_split(PolarizedAllSplitPartialDensity, not_core): - msg = "set LSEPK and LSEPB to .FALSE. in the INCAR file." - with pytest.raises(NotImplemented) as excinfo: - PolarizedAllSplitPartialDensity.to_stm(selection="constant_current") - assert msg in str(excinfo.value) - - -def test_to_stm_nonsplit_tip_to_high(NonSplitPartialDensity, not_core): - actual = NonSplitPartialDensity - tip_height = 8.4 - error = f"""The tip position at {tip_height:.2f} is above half of the - estimated vacuum thickness {actual._estimate_vacuum():.2f} Angstrom. - You would be sampling the bottom of your slab, which is not supported.""" - with pytest.raises(IncorrectUsage, match=error): - actual.to_stm(tip_height=tip_height) - - -def test_to_stm_nonsplit_not_orthogonal_no_vacuum( - PolarizedNonSplitPartialDensitySr2TiO4, - not_core, -): - msg = "The vacuum region in your cell is too small for STM simulations." - with pytest.raises(IncorrectUsage) as excinfo: - PolarizedNonSplitPartialDensitySr2TiO4.to_stm() - assert msg in str(excinfo.value) - - -def test_to_stm_wrong_spin_nonsplit(PolarizedNonSplitPartialDensity, not_core): - msg = "'up', 'down', or 'total'" - with pytest.raises(IncorrectUsage) as excinfo: - PolarizedNonSplitPartialDensity.to_stm(selection="all") - assert msg in str(excinfo.value) - - -def test_to_stm_wrong_mode(PolarizedNonSplitPartialDensity, not_core): - with pytest.raises(IncorrectUsage) as excinfo: - PolarizedNonSplitPartialDensity.to_stm(selection="stm") - assert "STM mode" in str(excinfo.value) - - -def test_wrong_vacuum_direction(NonSplitPartialDensityNi_100, not_core): - msg = """The vacuum region in your cell is not located along - the third lattice vector.""" - with pytest.raises(NotImplemented) as excinfo: - NonSplitPartialDensityNi_100.to_stm() - assert msg in str(excinfo.value) - - -@pytest.mark.parametrize("alias", ("constant_height", "ch", "height")) -def test_to_stm_nonsplit_constant_height( - PolarizedNonSplitPartialDensity, alias, spin, Assert, not_core -): - supercell = 3 - actual = PolarizedNonSplitPartialDensity.to_stm( - selection=f"{alias}({spin})", tip_height=2.0, supercell=supercell - ) - expected = PolarizedNonSplitPartialDensity.ref - assert type(actual.series.data) == np.ndarray - assert actual.series.data.shape == (expected.grid[0], expected.grid[1]) - Assert.allclose(actual.series.lattice.vectors, expected.plane_vectors.vectors) - Assert.allclose(actual.series.supercell, np.asarray([supercell, supercell])) - # check different elements of the label - assert type(actual.series.label) is str - expected = "both spin channels" if spin == "total" else f"spin {spin}" - assert expected in actual.series.label - assert "constant height" in actual.series.label - assert "2.0" in actual.series.label - assert "constant height" in actual.title - assert "2.0" in actual.title - - -@pytest.mark.parametrize("alias", ("constant_current", "cc", "current")) -def test_to_stm_nonsplit_constant_current( - PolarizedNonSplitPartialDensity, alias, spin, Assert, not_core -): - current = 5 - supercell = np.asarray([2, 4]) - actual = PolarizedNonSplitPartialDensity.to_stm( - selection=f"{spin}({alias})", - current=current, - supercell=supercell, - ) - expected = PolarizedNonSplitPartialDensity.ref - assert type(actual.series.data) == np.ndarray - assert actual.series.data.shape == (expected.grid[0], expected.grid[1]) - Assert.allclose(actual.series.lattice.vectors, expected.plane_vectors.vectors) - Assert.allclose(actual.series.supercell, supercell) - # check different elements of the label - assert type(actual.series.label) is str - expected = "both spin channels" if spin == "total" else f"spin {spin}" - assert expected in actual.series.label - assert "constant current" in actual.series.label - assert f"{current:.2f}" in actual.series.label - assert "constant current" in actual.title - assert f"{current:.2f}" in actual.title - - -@pytest.mark.parametrize("alias", ("constant_current", "cc", "current")) -def test_to_stm_nonsplit_constant_current_non_ortho( - NonSplitPartialDensityCaAs3_110, alias, spin, Assert, not_core -): - current = 5 - supercell = np.asarray([2, 4]) - actual = NonSplitPartialDensityCaAs3_110.to_stm( - selection=f"{spin}({alias})", - current=current, - supercell=supercell, - ) - expected = NonSplitPartialDensityCaAs3_110.ref - assert type(actual.series.data) == np.ndarray - assert actual.series.data.shape == (expected.grid[0], expected.grid[1]) - Assert.allclose(actual.series.lattice.vectors, expected.plane_vectors.vectors) - Assert.allclose(actual.series.supercell, supercell) - # check different elements of the label - assert type(actual.series.label) is str - expected = "both spin channels" if spin == "total" else f"spin {spin}" - assert expected in actual.series.label - assert "constant current" in actual.series.label - assert f"{current:.2f}" in actual.series.label - assert "constant current" in actual.title - assert f"{current:.2f}" in actual.title - - -def test_stm_default_settings(PolarizedNonSplitPartialDensity, not_core): - actual = dataclasses.asdict(PolarizedNonSplitPartialDensity.stm_settings) - defaults = { - "sigma_xy": 4.0, - "sigma_z": 4.0, - "truncate": 3.0, - "enhancement_factor": 1000, - "interpolation_factor": 10, - } - assert actual == defaults - modified = PartialDensity.STM_settings( - sigma_xy=2.0, - sigma_z=2.0, - truncate=1.0, - enhancement_factor=500, - interpolation_factor=5, - ) - graph = PolarizedNonSplitPartialDensity.to_stm(stm_settings=modified) - assert graph.series.settings == modified - - -def test_smoothening_change(PolarizedNonSplitPartialDensity, not_core): - mod_settings = PartialDensity.STM_settings(sigma_xy=2.0, sigma_z=2.0, truncate=1.0) - data = PolarizedNonSplitPartialDensity.to_numpy("total", band=0, kpoint=0) - default_smoothed_density = PolarizedNonSplitPartialDensity._smooth_stm_data( - data=data, stm_settings=PartialDensity.STM_settings() - ) - new_smoothed_density = PolarizedNonSplitPartialDensity._smooth_stm_data( - data=data, stm_settings=mod_settings - ) - assert not np.allclose(default_smoothed_density, new_smoothed_density) - - -def test_enhancement_setting_change(PolarizedNonSplitPartialDensity, Assert, not_core): - enhance_settings = PartialDensity.STM_settings( - enhancement_factor=PartialDensity.STM_settings().enhancement_factor / 2.0 - ) - graph_def = PolarizedNonSplitPartialDensity.to_stm("constant_height") - graph_less_enhanced = PolarizedNonSplitPartialDensity.to_stm( - "constant_height", stm_settings=enhance_settings - ) - Assert.allclose(graph_def.series.data, graph_less_enhanced.series.data * 2) - - -def test_interpolation_setting_change(PolarizedNonSplitPartialDensity, not_core): - interp_settings = PartialDensity.STM_settings( - interpolation_factor=PartialDensity.STM_settings().interpolation_factor / 4.0 - ) - graph_def = PolarizedNonSplitPartialDensity.to_stm("constant_current", current=1) - graph_less_interp_points = PolarizedNonSplitPartialDensity.to_stm( - "constant_current", current=1, stm_settings=interp_settings - ) - assert not np.allclose(graph_def.series.data, graph_less_interp_points.series.data) - - -def test_factory_methods(raw_data, check_factory_methods, not_core): - data = raw_data.partial_density("spin_polarized") - check_factory_methods(PartialDensity, data) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import dataclasses +import types + +import numpy as np +import pytest + +from py4vasp import _config +from py4vasp._calculation.partial_density import PartialDensity +from py4vasp._calculation.structure import Structure, StructureHandler +from py4vasp._util.slicing import plane +from py4vasp.exception import IncorrectUsage, NoData, NotImplemented + + +@pytest.fixture( + params=[ + "no splitting no spin", + "no splitting no spin Ca3AsBr3", + "no splitting no spin Sr2TiO4", + "no splitting no spin CaAs3_110", + "split_bands", + "split_bands and spin_polarized", + "split_bands and spin_polarized Ca3AsBr3", + "split_bands and spin_polarized Sr2TiO4", + "split_kpoints", + "split_kpoints and spin_polarized", + "split_kpoints and spin_polarized Ca3AsBr3", + "split_kpoints and spin_polarized Sr2TiO4", + "spin_polarized", + "spin_polarized Ca3AsBr3", + "spin_polarized Sr2TiO4", + "split_bands and split_kpoints", + "split_bands and split_kpoints and spin_polarized", + "split_bands and split_kpoints and spin_polarized Ca3AsBr3", + "split_bands and split_kpoints and spin_polarized Sr2TiO4", + ] +) +def AnyPartialDensity(raw_data, request): + return make_reference_partial_density(raw_data, request.param) + + +@pytest.fixture +def NonSplitPartialDensity(raw_data): + return make_reference_partial_density(raw_data, "no splitting no spin") + + +@pytest.fixture +def PolarizedNonSplitPartialDensity(raw_data): + return make_reference_partial_density(raw_data, "spin_polarized") + + +@pytest.fixture +def PolarizedNonSplitPartialDensityCa3AsBr3(raw_data): + return make_reference_partial_density(raw_data, "spin_polarized Ca3AsBr3") + + +@pytest.fixture +def NonSplitPartialDensityCaAs3_110(raw_data): + return make_reference_partial_density(raw_data, "CaAs3_110") + + +@pytest.fixture +def NonSplitPartialDensityNi_100(raw_data): + return make_reference_partial_density(raw_data, "Ni100") + + +@pytest.fixture +def PolarizedNonSplitPartialDensitySr2TiO4(raw_data): + return make_reference_partial_density(raw_data, "spin_polarized Sr2TiO4") + + +@pytest.fixture +def NonPolarizedBandSplitPartialDensity(raw_data): + return make_reference_partial_density(raw_data, "split_bands") + + +@pytest.fixture +def PolarizedAllSplitPartialDensity(raw_data): + return make_reference_partial_density( + raw_data, "split_bands and split_kpoints and spin_polarized" + ) + + +@pytest.fixture(params=["up", "down", "total"]) +def spin(request): + return request.param + + +def make_reference_partial_density(raw_data, selection): + raw_partial_density = raw_data.partial_density(selection=selection) + parchg = PartialDensity.from_data(raw_partial_density) + parchg.ref = types.SimpleNamespace() + parchg.ref.structure = Structure.from_data(raw_partial_density.structure) + parchg.ref.structure_handler = StructureHandler.from_data( + raw_partial_density.structure + ) + parchg.ref.plane_vectors = plane( + cell=parchg.ref.structure.lattice_vectors(), + cut="c", + normal="z", + ) + parchg.ref.partial_density = raw_partial_density.partial_charge + parchg.ref.bands = raw_partial_density.bands + parchg.ref.kpoints = raw_partial_density.kpoints + parchg.ref.grid = raw_partial_density.grid + return parchg + + +def test_read(AnyPartialDensity, Assert, not_core): + actual = AnyPartialDensity.read() + expected = AnyPartialDensity.ref + Assert.allclose(actual["bands"], expected.bands) + Assert.allclose(actual["kpoints"], expected.kpoints) + Assert.allclose(actual["grid"], expected.grid) + expected_density = np.squeeze(np.asarray(expected.partial_density).T) + Assert.allclose(actual["partial_density"], expected_density) + Assert.same_structure(actual["structure"], expected.structure.read()) + + +def test_stoichiometry(AnyPartialDensity, not_core): + actual = AnyPartialDensity._stoichiometry() + expected = str(AnyPartialDensity.ref.structure_handler._stoichiometry()) + assert actual == expected + + +def test_bands(AnyPartialDensity, Assert, not_core): + actual = AnyPartialDensity.bands() + expected = AnyPartialDensity.ref.bands + Assert.allclose(actual, expected) + + +def test_kpoints(AnyPartialDensity, Assert, not_core): + actual = AnyPartialDensity.kpoints() + expected = AnyPartialDensity.ref.kpoints + Assert.allclose(actual, expected) + + +def test_grid(AnyPartialDensity, Assert, not_core): + actual = AnyPartialDensity.grid() + expected = AnyPartialDensity.ref.grid + Assert.allclose(actual, expected) + + +@pytest.mark.parametrize("selection", (None, "total", "up", "down")) +def test_plot(PolarizedNonSplitPartialDensity, selection, Assert, not_core): + if selection is not None: + view = PolarizedNonSplitPartialDensity.plot(selection) + expected_density = PolarizedNonSplitPartialDensity.to_numpy(selection) + expected_label = selection + else: + view = PolarizedNonSplitPartialDensity.plot() + expected_density = PolarizedNonSplitPartialDensity.to_numpy() + expected_label = "total" + reference = PolarizedNonSplitPartialDensity.ref + Assert.same_structure_view(view, reference.structure.plot()) + assert len(view.grid_scalars) == 1 + grid_scalar = view.grid_scalars[0] + assert grid_scalar.label == expected_label + assert grid_scalar.quantity.ndim == 4 + Assert.allclose(grid_scalar.quantity, expected_density) + assert len(grid_scalar.isosurfaces) == 1 + isosurface = grid_scalar.isosurfaces[0] + assert isosurface.isolevel == 0.2 + assert isosurface.color == _config.VASP_COLORS["cyan"] + assert isosurface.opacity == 0.6 + + +def test_plot_with_options(NonSplitPartialDensity, Assert, not_core): + view = NonSplitPartialDensity.plot(supercell=(1, 2, 3), isolevel=0.6, color="red") + Assert.allclose(view.supercell, (1, 2, 3)) + assert len(view.grid_scalars) == 1 + grid_scalar = view.grid_scalars[0] + assert len(grid_scalar.isosurfaces) == 1 + isosurface = grid_scalar.isosurfaces[0] + assert isosurface.isolevel == 0.6 + assert isosurface.color == "red" + + +def test_non_split_to_numpy(PolarizedNonSplitPartialDensity, Assert, not_core): + actual = PolarizedNonSplitPartialDensity.to_numpy("total") + expected = PolarizedNonSplitPartialDensity.ref.partial_density + Assert.allclose(actual, expected[0, 0, 0].T) + + actual = PolarizedNonSplitPartialDensity.to_numpy("up") + Assert.allclose(actual, 0.5 * (expected[0, 0, 0].T + expected[0, 0, 1].T)) + + actual = PolarizedNonSplitPartialDensity.to_numpy("down") + Assert.allclose(actual, 0.5 * (expected[0, 0, 0].T - expected[0, 0, 1].T)) + + +def test_split_to_numpy(PolarizedAllSplitPartialDensity, Assert, not_core): + bands = PolarizedAllSplitPartialDensity.ref.bands + kpoints = PolarizedAllSplitPartialDensity.ref.kpoints + for band_index, band in enumerate(bands): + for kpoint_index, kpoint in enumerate(kpoints): + actual = PolarizedAllSplitPartialDensity.to_numpy( + band=band, kpoint=kpoint, selection="total" + ) + expected = PolarizedAllSplitPartialDensity.ref.partial_density + Assert.allclose(actual, np.asarray(expected)[kpoint_index, band_index, 0].T) + + msg = f"Band {max(bands) + 1} not found in the bands array." + with pytest.raises(NoData) as excinfo: + PolarizedAllSplitPartialDensity.to_numpy( + band=max(bands) + 1, kpoint=max(kpoints), selection="up" + ) + assert msg in str(excinfo.value) + + msg = f"K-point {min(kpoints) - 1} not found in the kpoints array." + with pytest.raises(NoData) as excinfo: + PolarizedAllSplitPartialDensity.to_numpy( + band=min(bands), kpoint=min(kpoints) - 1, selection="down" + ) + assert msg in str(excinfo.value) + + +def test_non_polarized_to_numpy(NonSplitPartialDensity, spin, Assert, not_core): + actual = NonSplitPartialDensity.to_numpy(selection=spin) + expected = NonSplitPartialDensity.ref.partial_density + Assert.allclose(actual, np.asarray(expected).T[:, :, :, 0, 0, 0]) + + +def test_split_bands_to_numpy( + NonPolarizedBandSplitPartialDensity, spin, Assert, not_core +): + bands = NonPolarizedBandSplitPartialDensity.ref.bands + for band_index, band in enumerate(bands): + actual = NonPolarizedBandSplitPartialDensity.to_numpy(spin, band=band) + expected = NonPolarizedBandSplitPartialDensity.ref.partial_density + Assert.allclose(actual, np.asarray(expected).T[:, :, :, 0, band_index, 0]) + + +def test_to_stm_split(PolarizedAllSplitPartialDensity, not_core): + msg = "set LSEPK and LSEPB to .FALSE. in the INCAR file." + with pytest.raises(NotImplemented) as excinfo: + PolarizedAllSplitPartialDensity.to_stm(selection="constant_current") + assert msg in str(excinfo.value) + + +def test_to_stm_nonsplit_tip_to_high(NonSplitPartialDensity, not_core): + actual = NonSplitPartialDensity + tip_height = 8.4 + error = f"""The tip position at {tip_height:.2f} is above half of the + estimated vacuum thickness {actual._estimate_vacuum():.2f} Angstrom. + You would be sampling the bottom of your slab, which is not supported.""" + with pytest.raises(IncorrectUsage, match=error): + actual.to_stm(tip_height=tip_height) + + +def test_to_stm_nonsplit_not_orthogonal_no_vacuum( + PolarizedNonSplitPartialDensitySr2TiO4, + not_core, +): + msg = "The vacuum region in your cell is too small for STM simulations." + with pytest.raises(IncorrectUsage) as excinfo: + PolarizedNonSplitPartialDensitySr2TiO4.to_stm() + assert msg in str(excinfo.value) + + +def test_to_stm_wrong_spin_nonsplit(PolarizedNonSplitPartialDensity, not_core): + msg = "'up', 'down', or 'total'" + with pytest.raises(IncorrectUsage) as excinfo: + PolarizedNonSplitPartialDensity.to_stm(selection="all") + assert msg in str(excinfo.value) + + +def test_to_stm_wrong_mode(PolarizedNonSplitPartialDensity, not_core): + with pytest.raises(IncorrectUsage) as excinfo: + PolarizedNonSplitPartialDensity.to_stm(selection="stm") + assert "STM mode" in str(excinfo.value) + + +def test_wrong_vacuum_direction(NonSplitPartialDensityNi_100, not_core): + msg = """The vacuum region in your cell is not located along + the third lattice vector.""" + with pytest.raises(NotImplemented) as excinfo: + NonSplitPartialDensityNi_100.to_stm() + assert msg in str(excinfo.value) + + +@pytest.mark.parametrize("alias", ("constant_height", "ch", "height")) +def test_to_stm_nonsplit_constant_height( + PolarizedNonSplitPartialDensity, alias, spin, Assert, not_core +): + supercell = 3 + actual = PolarizedNonSplitPartialDensity.to_stm( + selection=f"{alias}({spin})", tip_height=2.0, supercell=supercell + ) + expected = PolarizedNonSplitPartialDensity.ref + assert type(actual.series.data) == np.ndarray + assert actual.series.data.shape == (expected.grid[0], expected.grid[1]) + Assert.allclose(actual.series.lattice.vectors, expected.plane_vectors.vectors) + Assert.allclose(actual.series.supercell, np.asarray([supercell, supercell])) + # check different elements of the label + assert type(actual.series.label) is str + expected = "both spin channels" if spin == "total" else f"spin {spin}" + assert expected in actual.series.label + assert "constant height" in actual.series.label + assert "2.0" in actual.series.label + assert "constant height" in actual.title + assert "2.0" in actual.title + + +@pytest.mark.parametrize("alias", ("constant_current", "cc", "current")) +def test_to_stm_nonsplit_constant_current( + PolarizedNonSplitPartialDensity, alias, spin, Assert, not_core +): + current = 5 + supercell = np.asarray([2, 4]) + actual = PolarizedNonSplitPartialDensity.to_stm( + selection=f"{spin}({alias})", + current=current, + supercell=supercell, + ) + expected = PolarizedNonSplitPartialDensity.ref + assert type(actual.series.data) == np.ndarray + assert actual.series.data.shape == (expected.grid[0], expected.grid[1]) + Assert.allclose(actual.series.lattice.vectors, expected.plane_vectors.vectors) + Assert.allclose(actual.series.supercell, supercell) + # check different elements of the label + assert type(actual.series.label) is str + expected = "both spin channels" if spin == "total" else f"spin {spin}" + assert expected in actual.series.label + assert "constant current" in actual.series.label + assert f"{current:.2f}" in actual.series.label + assert "constant current" in actual.title + assert f"{current:.2f}" in actual.title + + +@pytest.mark.parametrize("alias", ("constant_current", "cc", "current")) +def test_to_stm_nonsplit_constant_current_non_ortho( + NonSplitPartialDensityCaAs3_110, alias, spin, Assert, not_core +): + current = 5 + supercell = np.asarray([2, 4]) + actual = NonSplitPartialDensityCaAs3_110.to_stm( + selection=f"{spin}({alias})", + current=current, + supercell=supercell, + ) + expected = NonSplitPartialDensityCaAs3_110.ref + assert type(actual.series.data) == np.ndarray + assert actual.series.data.shape == (expected.grid[0], expected.grid[1]) + Assert.allclose(actual.series.lattice.vectors, expected.plane_vectors.vectors) + Assert.allclose(actual.series.supercell, supercell) + # check different elements of the label + assert type(actual.series.label) is str + expected = "both spin channels" if spin == "total" else f"spin {spin}" + assert expected in actual.series.label + assert "constant current" in actual.series.label + assert f"{current:.2f}" in actual.series.label + assert "constant current" in actual.title + assert f"{current:.2f}" in actual.title + + +def test_stm_default_settings(PolarizedNonSplitPartialDensity, not_core): + actual = dataclasses.asdict(PolarizedNonSplitPartialDensity.stm_settings) + defaults = { + "sigma_xy": 4.0, + "sigma_z": 4.0, + "truncate": 3.0, + "enhancement_factor": 1000, + "interpolation_factor": 10, + } + assert actual == defaults + modified = PartialDensity.STM_settings( + sigma_xy=2.0, + sigma_z=2.0, + truncate=1.0, + enhancement_factor=500, + interpolation_factor=5, + ) + graph = PolarizedNonSplitPartialDensity.to_stm(stm_settings=modified) + assert graph.series.settings == modified + + +def test_smoothening_change(PolarizedNonSplitPartialDensity, not_core): + mod_settings = PartialDensity.STM_settings(sigma_xy=2.0, sigma_z=2.0, truncate=1.0) + data = PolarizedNonSplitPartialDensity.to_numpy("total", band=0, kpoint=0) + default_smoothed_density = PolarizedNonSplitPartialDensity._smooth_stm_data( + data=data, stm_settings=PartialDensity.STM_settings() + ) + new_smoothed_density = PolarizedNonSplitPartialDensity._smooth_stm_data( + data=data, stm_settings=mod_settings + ) + assert not np.allclose(default_smoothed_density, new_smoothed_density) + + +def test_enhancement_setting_change(PolarizedNonSplitPartialDensity, Assert, not_core): + enhance_settings = PartialDensity.STM_settings( + enhancement_factor=PartialDensity.STM_settings().enhancement_factor / 2.0 + ) + graph_def = PolarizedNonSplitPartialDensity.to_stm("constant_height") + graph_less_enhanced = PolarizedNonSplitPartialDensity.to_stm( + "constant_height", stm_settings=enhance_settings + ) + Assert.allclose(graph_def.series.data, graph_less_enhanced.series.data * 2) + + +def test_interpolation_setting_change(PolarizedNonSplitPartialDensity, not_core): + interp_settings = PartialDensity.STM_settings( + interpolation_factor=PartialDensity.STM_settings().interpolation_factor / 4.0 + ) + graph_def = PolarizedNonSplitPartialDensity.to_stm("constant_current", current=1) + graph_less_interp_points = PolarizedNonSplitPartialDensity.to_stm( + "constant_current", current=1, stm_settings=interp_settings + ) + assert not np.allclose(graph_def.series.data, graph_less_interp_points.series.data) + + +def test_factory_methods(raw_data, check_factory_methods, not_core): + data = raw_data.partial_density("spin_polarized") + check_factory_methods(PartialDensity, data) diff --git a/tests/calculation/test_selection_convention.py b/tests/calculation/test_selection_convention.py index d26f1bb0e..0456e242c 100644 --- a/tests/calculation/test_selection_convention.py +++ b/tests/calculation/test_selection_convention.py @@ -1,303 +1,303 @@ -# Copyright © VASP Software GmbH, -# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -"""Test that all dispatcher methods follow the selection parameter convention. - -The dispatch system uses `merge_X(source, quantity_name, selection, handler_factory, -method, *args, **kwargs)` where: -- The 3rd argument (`selection`) is used for source routing AND is automatically - forwarded (as `remaining_selection`) to the handler method when it accepts a - `selection` parameter. -- `*args` are extra arguments forwarded to the handler method AFTER the automatic - selection injection. Selection must NEVER appear in *args. - -Rules: -1. If the dispatcher method has a `selection` parameter: - - The merge call's 3rd argument MUST be `selection`. -2. If the dispatcher method does NOT have a `selection` parameter: - - The merge call's 3rd argument MUST be `None`. -3. In ALL cases: `selection` must NOT appear in *args (the dispatch system - handles forwarding automatically). -""" - -import ast -import importlib -import inspect -import pathlib - -import pytest - -from py4vasp._calculation.dispatch import _REGISTRY -from py4vasp._raw.definition import schema - -# Force-import all dispatcher modules so the registry is populated. -_CALCULATION_DIR = ( - pathlib.Path(__file__).resolve().parent.parent.parent - / "src" - / "py4vasp" - / "_calculation" -) -for _f in sorted(_CALCULATION_DIR.glob("*.py")): - if _f.name.startswith("_") and _f.name != "_CONTCAR.py": - continue - _module_name = _f.stem - try: - importlib.import_module(f"py4vasp._calculation.{_module_name}") - except (ImportError, Exception): - pass - -MERGE_FUNCS = frozenset({"merge_default", "merge_graphs", "merge_strings"}) - -# Decorator quantity names that correspond to multi-source schema entries. -# Built from the schema at import time. -MULTI_SOURCE_NAMES = frozenset( - qty for qty, srcs in schema._sources.items() if len(srcs) > 1 -) - -# Some @quantity decorators use names that differ from the schema key. -# Map decorator name -> schema name for multi-source lookup. -_DECORATOR_TO_SCHEMA = { - "transport": "electron_phonon_transport", -} - -# Classes that use a non-standard selection pattern (e.g. self._selection_name) -# and should be excluded from the standard convention check. -_EXCLUDED_CLASSES = frozenset({"Density"}) - - -def _get_all_dispatcher_classes(): - """Yield (quantity_name, cls) for every registered dispatcher class.""" - for key, value in _REGISTRY.items(): - if isinstance(value, dict): - # Group (e.g. phonon -> {band: cls, dos: cls}) - for sub_name, cls in value.items(): - yield sub_name, cls - else: - yield key, cls - - -def _get_source_file(cls): - """Return the Path to the source file of a class.""" - return pathlib.Path(inspect.getfile(cls)) - - -def _parse_class_ast(cls): - """Return the AST ClassDef node for the given class.""" - source_file = _get_source_file(cls) - tree = ast.parse(source_file.read_text(encoding="utf-8")) - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and node.name == cls.__name__: - return node, tree - raise ValueError(f"Could not find class {cls.__name__} in {source_file}") - - -def _get_handler_classes_from_file(tree): - """Return a dict of {ClassName: ClassDef} for non-dispatcher classes.""" - handlers = {} - for node in ast.walk(tree): - if not isinstance(node, ast.ClassDef): - continue - is_dispatcher = any( - isinstance(d, ast.Call) - and isinstance(d.func, ast.Name) - and d.func.id == "quantity" - for d in node.decorator_list - ) - if not is_dispatcher: - handlers[node.name] = node - return handlers - - -def _get_handler_method_params(handler_classes, handler_class_name, method_name): - """Return the list of parameter names for a handler method.""" - cls_node = handler_classes.get(handler_class_name) - if cls_node is None: - return [] - for item in cls_node.body: - if isinstance(item, ast.FunctionDef) and item.name == method_name: - return [arg.arg for arg in item.args.args] - return [] - - -def _get_ast_repr(node): - """Get a string representation of an AST node for comparison.""" - if isinstance(node, ast.Constant) and node.value is None: - return "None" - elif isinstance(node, ast.Name): - return node.id - elif isinstance(node, ast.Attribute): - parts = [] - n = node - while isinstance(n, ast.Attribute): - parts.append(n.attr) - n = n.value - if isinstance(n, ast.Name): - parts.append(n.id) - return ".".join(reversed(parts)) - return "other" - - -def _find_merge_calls(method_node): - """Find all merge_X calls within a method and return analysis info.""" - calls = [] - for node in ast.walk(method_node): - if not isinstance(node, ast.Call): - continue - func_name = getattr(node.func, "id", None) - if func_name not in MERGE_FUNCS: - continue - # Extract 3rd argument (selection for source routing) - third_arg = _get_ast_repr(node.args[2]) if len(node.args) >= 3 else "missing" - # Extract 5th argument (handler method reference) - handler_method_ref = None - if len(node.args) >= 5: - ref = node.args[4] - if isinstance(ref, ast.Attribute) and isinstance(ref.value, ast.Name): - handler_method_ref = (ref.value.id, ref.attr) - # Extract extra args (after handler method ref, i.e. args[5:]) - extra_args = [_get_ast_repr(a) for a in node.args[5:]] - calls.append( - { - "func": func_name, - "third_arg": third_arg, - "handler_ref": handler_method_ref, - "extra_args": extra_args, - } - ) - return calls - - -def _is_public_method(method_node): - """Check if a method is public (not starting with _) and not a dunder helper.""" - name = method_node.name - if name.startswith("_") and not name.startswith("__"): - return False - # Skip non-dispatch helpers - if name in ("__init__", "__getitem__", "__copy__", "_repr_pretty_"): - return False - return True - - -def _get_quantity_name_from_decorator(cls_node): - """Extract the quantity name string from @quantity('name') decorator.""" - for d in cls_node.decorator_list: - if ( - isinstance(d, ast.Call) - and isinstance(d.func, ast.Name) - and d.func.id == "quantity" - ): - if d.args and isinstance(d.args[0], ast.Constant): - return d.args[0].value - return None - - -def _collect_test_cases(): - """Collect all (quantity_key, class, method_name, merge_call) test cases.""" - cases = [] - for key, value in _REGISTRY.items(): - if isinstance(value, dict): - for sub_key, cls in value.items(): - cases.extend(_cases_for_class(sub_key, cls)) - else: - cases.extend(_cases_for_class(key, value)) - return cases - - -def _cases_for_class(registry_key, cls): - """Generate test cases for a single dispatcher class.""" - cases = [] - if cls.__name__ in _EXCLUDED_CLASSES: - return cases - try: - cls_node, tree = _parse_class_ast(cls) - except (ValueError, OSError): - return cases - - qty_name = _get_quantity_name_from_decorator(cls_node) - if qty_name is None: - return cases - - # Map decorator name to schema name for multi-source check - schema_name = _DECORATOR_TO_SCHEMA.get(qty_name, qty_name) - is_multi = schema_name in MULTI_SOURCE_NAMES - handler_classes = _get_handler_classes_from_file(tree) - - for item in cls_node.body: - if not isinstance(item, ast.FunctionDef): - continue - if not _is_public_method(item): - continue - - merge_calls = _find_merge_calls(item) - if not merge_calls: - continue - - method_params = [arg.arg for arg in item.args.args] - has_selection_param = "selection" in method_params - - for call_info in merge_calls: - handler_ref = call_info["handler_ref"] - handler_has_selection = False - if handler_ref: - handler_class_name, handler_method_name = handler_ref - handler_params = _get_handler_method_params( - handler_classes, handler_class_name, handler_method_name - ) - handler_has_selection = "selection" in handler_params - - cases.append( - ( - registry_key, - cls.__name__, - item.name, - is_multi, - handler_has_selection, - has_selection_param, - call_info, - ) - ) - return cases - - -_ALL_CASES = _collect_test_cases() - - -def _case_id(case): - registry_key, cls_name, method_name, *_ = case - return f"{cls_name}.{method_name}" - - -@pytest.mark.parametrize("case", _ALL_CASES, ids=_case_id) -def test_selection_convention(case): - ( - registry_key, - cls_name, - method_name, - is_multi, - handler_has_selection, - has_selection_param, - call_info, - ) = case - - third_arg = call_info["third_arg"] - extra_args = call_info["extra_args"] - - # Rule 1: If the dispatcher has a `selection` parameter, the 3rd merge arg - # MUST be `selection` (enables source routing AND auto-forwarding). - # Rule 2: If it does NOT have `selection`, the 3rd arg MUST be `None`. - if has_selection_param: - assert third_arg == "selection", ( - f"{cls_name}.{method_name}: dispatcher has `selection` parameter, " - f"so 3rd merge argument must be `selection`, got `{third_arg}`" - ) - else: - assert third_arg == "None", ( - f"{cls_name}.{method_name}: dispatcher has no `selection` parameter, " - f"so 3rd merge argument must be `None`, got `{third_arg}`" - ) - - # Rule 3: `selection` must NEVER appear in *args — the dispatch system - # handles forwarding automatically via introspection. - assert "selection" not in extra_args, ( - f"{cls_name}.{method_name}: `selection` must not be passed in *args " - f"(dispatch auto-forwards it). Extra args: {extra_args}" - ) +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +"""Test that all dispatcher methods follow the selection parameter convention. + +The dispatch system uses `merge_X(source, quantity_name, selection, handler_factory, +method, *args, **kwargs)` where: +- The 3rd argument (`selection`) is used for source routing AND is automatically + forwarded (as `remaining_selection`) to the handler method when it accepts a + `selection` parameter. +- `*args` are extra arguments forwarded to the handler method AFTER the automatic + selection injection. Selection must NEVER appear in *args. + +Rules: +1. If the dispatcher method has a `selection` parameter: + - The merge call's 3rd argument MUST be `selection`. +2. If the dispatcher method does NOT have a `selection` parameter: + - The merge call's 3rd argument MUST be `None`. +3. In ALL cases: `selection` must NOT appear in *args (the dispatch system + handles forwarding automatically). +""" + +import ast +import importlib +import inspect +import pathlib + +import pytest + +from py4vasp._calculation.dispatch import _REGISTRY +from py4vasp._raw.definition import schema + +# Force-import all dispatcher modules so the registry is populated. +_CALCULATION_DIR = ( + pathlib.Path(__file__).resolve().parent.parent.parent + / "src" + / "py4vasp" + / "_calculation" +) +for _f in sorted(_CALCULATION_DIR.glob("*.py")): + if _f.name.startswith("_") and _f.name != "_CONTCAR.py": + continue + _module_name = _f.stem + try: + importlib.import_module(f"py4vasp._calculation.{_module_name}") + except (ImportError, Exception): + pass + +MERGE_FUNCS = frozenset({"merge_default", "merge_graphs", "merge_strings"}) + +# Decorator quantity names that correspond to multi-source schema entries. +# Built from the schema at import time. +MULTI_SOURCE_NAMES = frozenset( + qty for qty, srcs in schema._sources.items() if len(srcs) > 1 +) + +# Some @quantity decorators use names that differ from the schema key. +# Map decorator name -> schema name for multi-source lookup. +_DECORATOR_TO_SCHEMA = { + "transport": "electron_phonon_transport", +} + +# Classes that use a non-standard selection pattern (e.g. self._selection_name) +# and should be excluded from the standard convention check. +_EXCLUDED_CLASSES = frozenset({"Density"}) + + +def _get_all_dispatcher_classes(): + """Yield (quantity_name, cls) for every registered dispatcher class.""" + for key, value in _REGISTRY.items(): + if isinstance(value, dict): + # Group (e.g. phonon -> {band: cls, dos: cls}) + for sub_name, cls in value.items(): + yield sub_name, cls + else: + yield key, cls + + +def _get_source_file(cls): + """Return the Path to the source file of a class.""" + return pathlib.Path(inspect.getfile(cls)) + + +def _parse_class_ast(cls): + """Return the AST ClassDef node for the given class.""" + source_file = _get_source_file(cls) + tree = ast.parse(source_file.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == cls.__name__: + return node, tree + raise ValueError(f"Could not find class {cls.__name__} in {source_file}") + + +def _get_handler_classes_from_file(tree): + """Return a dict of {ClassName: ClassDef} for non-dispatcher classes.""" + handlers = {} + for node in ast.walk(tree): + if not isinstance(node, ast.ClassDef): + continue + is_dispatcher = any( + isinstance(d, ast.Call) + and isinstance(d.func, ast.Name) + and d.func.id == "quantity" + for d in node.decorator_list + ) + if not is_dispatcher: + handlers[node.name] = node + return handlers + + +def _get_handler_method_params(handler_classes, handler_class_name, method_name): + """Return the list of parameter names for a handler method.""" + cls_node = handler_classes.get(handler_class_name) + if cls_node is None: + return [] + for item in cls_node.body: + if isinstance(item, ast.FunctionDef) and item.name == method_name: + return [arg.arg for arg in item.args.args] + return [] + + +def _get_ast_repr(node): + """Get a string representation of an AST node for comparison.""" + if isinstance(node, ast.Constant) and node.value is None: + return "None" + elif isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Attribute): + parts = [] + n = node + while isinstance(n, ast.Attribute): + parts.append(n.attr) + n = n.value + if isinstance(n, ast.Name): + parts.append(n.id) + return ".".join(reversed(parts)) + return "other" + + +def _find_merge_calls(method_node): + """Find all merge_X calls within a method and return analysis info.""" + calls = [] + for node in ast.walk(method_node): + if not isinstance(node, ast.Call): + continue + func_name = getattr(node.func, "id", None) + if func_name not in MERGE_FUNCS: + continue + # Extract 3rd argument (selection for source routing) + third_arg = _get_ast_repr(node.args[2]) if len(node.args) >= 3 else "missing" + # Extract 5th argument (handler method reference) + handler_method_ref = None + if len(node.args) >= 5: + ref = node.args[4] + if isinstance(ref, ast.Attribute) and isinstance(ref.value, ast.Name): + handler_method_ref = (ref.value.id, ref.attr) + # Extract extra args (after handler method ref, i.e. args[5:]) + extra_args = [_get_ast_repr(a) for a in node.args[5:]] + calls.append( + { + "func": func_name, + "third_arg": third_arg, + "handler_ref": handler_method_ref, + "extra_args": extra_args, + } + ) + return calls + + +def _is_public_method(method_node): + """Check if a method is public (not starting with _) and not a dunder helper.""" + name = method_node.name + if name.startswith("_") and not name.startswith("__"): + return False + # Skip non-dispatch helpers + if name in ("__init__", "__getitem__", "__copy__", "_repr_pretty_"): + return False + return True + + +def _get_quantity_name_from_decorator(cls_node): + """Extract the quantity name string from @quantity('name') decorator.""" + for d in cls_node.decorator_list: + if ( + isinstance(d, ast.Call) + and isinstance(d.func, ast.Name) + and d.func.id == "quantity" + ): + if d.args and isinstance(d.args[0], ast.Constant): + return d.args[0].value + return None + + +def _collect_test_cases(): + """Collect all (quantity_key, class, method_name, merge_call) test cases.""" + cases = [] + for key, value in _REGISTRY.items(): + if isinstance(value, dict): + for sub_key, cls in value.items(): + cases.extend(_cases_for_class(sub_key, cls)) + else: + cases.extend(_cases_for_class(key, value)) + return cases + + +def _cases_for_class(registry_key, cls): + """Generate test cases for a single dispatcher class.""" + cases = [] + if cls.__name__ in _EXCLUDED_CLASSES: + return cases + try: + cls_node, tree = _parse_class_ast(cls) + except (ValueError, OSError): + return cases + + qty_name = _get_quantity_name_from_decorator(cls_node) + if qty_name is None: + return cases + + # Map decorator name to schema name for multi-source check + schema_name = _DECORATOR_TO_SCHEMA.get(qty_name, qty_name) + is_multi = schema_name in MULTI_SOURCE_NAMES + handler_classes = _get_handler_classes_from_file(tree) + + for item in cls_node.body: + if not isinstance(item, ast.FunctionDef): + continue + if not _is_public_method(item): + continue + + merge_calls = _find_merge_calls(item) + if not merge_calls: + continue + + method_params = [arg.arg for arg in item.args.args] + has_selection_param = "selection" in method_params + + for call_info in merge_calls: + handler_ref = call_info["handler_ref"] + handler_has_selection = False + if handler_ref: + handler_class_name, handler_method_name = handler_ref + handler_params = _get_handler_method_params( + handler_classes, handler_class_name, handler_method_name + ) + handler_has_selection = "selection" in handler_params + + cases.append( + ( + registry_key, + cls.__name__, + item.name, + is_multi, + handler_has_selection, + has_selection_param, + call_info, + ) + ) + return cases + + +_ALL_CASES = _collect_test_cases() + + +def _case_id(case): + registry_key, cls_name, method_name, *_ = case + return f"{cls_name}.{method_name}" + + +@pytest.mark.parametrize("case", _ALL_CASES, ids=_case_id) +def test_selection_convention(case): + ( + registry_key, + cls_name, + method_name, + is_multi, + handler_has_selection, + has_selection_param, + call_info, + ) = case + + third_arg = call_info["third_arg"] + extra_args = call_info["extra_args"] + + # Rule 1: If the dispatcher has a `selection` parameter, the 3rd merge arg + # MUST be `selection` (enables source routing AND auto-forwarding). + # Rule 2: If it does NOT have `selection`, the 3rd arg MUST be `None`. + if has_selection_param: + assert third_arg == "selection", ( + f"{cls_name}.{method_name}: dispatcher has `selection` parameter, " + f"so 3rd merge argument must be `selection`, got `{third_arg}`" + ) + else: + assert third_arg == "None", ( + f"{cls_name}.{method_name}: dispatcher has no `selection` parameter, " + f"so 3rd merge argument must be `None`, got `{third_arg}`" + ) + + # Rule 3: `selection` must NEVER appear in *args — the dispatch system + # handles forwarding automatically via introspection. + assert "selection" not in extra_args, ( + f"{cls_name}.{method_name}: `selection` must not be passed in *args " + f"(dispatch auto-forwards it). Extra args: {extra_args}" + ) diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py index 8aa687436..c73967bf7 100644 --- a/tests/third_party/graph/test_graph.py +++ b/tests/third_party/graph/test_graph.py @@ -437,6 +437,16 @@ def test_merging_of_fields_of_graph(sine, parabola): dataclasses.replace(graph2, **{field.name: "other"}) + graph1 +def test_merging_of_numpy_field_values(parabola): + graph1 = Graph(parabola, xrange=np.array([0.0, 1.0])) + graph2 = Graph(parabola, xrange=np.array([0.0, 1.0])) + merged = graph1 + graph2 + assert np.array_equal(merged.xrange, np.array([0.0, 1.0])) + + with pytest.raises(exception.IncorrectUsage): + graph1 + Graph(parabola, xrange=np.array([0.0, 2.0])) + + def test_subplot(subplot, not_core): graph = Graph(subplot) graph.xlabel = ("first x-axis", "second x-axis") diff --git a/tests/third_party/view/test_vaspviewer.py b/tests/third_party/view/test_vaspviewer.py index b64399cd1..86957f65f 100644 --- a/tests/third_party/view/test_vaspviewer.py +++ b/tests/third_party/view/test_vaspviewer.py @@ -15,19 +15,33 @@ from py4vasp._third_party.view.view import GridQuantity, IonArrow, Isosurface from py4vasp._util import convert, import_ -vaspview = import_.optional("vasp_viewer") +vaspview = import_.optional("vasp.viewer") hasVaspView = pytest.mark.skipif( not import_.is_imported(vaspview), reason="vasp_viewer not installed", ) +def _trajectory_from_state(state_dict): + """Extract a numpy array from a viewer state dict with 'shape' and 'data' keys.""" + return np.frombuffer(bytes(state_dict["data"]), dtype=np.float32).reshape( + state_dict["shape"] + ) + + +def _quantity_from_state(state_dict): + """Extract a numpy array from a viewer state dict with 'quantity_shape' and 'quantity' keys.""" + return np.frombuffer(bytes(state_dict["quantity"]), dtype=np.float32).reshape( + state_dict["quantity_shape"] + ) + + def base_input_view(is_structure): if is_structure: return { - "atoms_types": [["Sr", "Ti", "O", "O", "O"]], + "elements": [["Sr", "Ti", "O", "O", "O"]], "lattice_vectors": [4 * np.eye(3)], - "atoms_trajectory": [ + "positions": [ [ [0.0, 0.0, 0.0], [0.5, 0.5, 0.5], @@ -39,12 +53,12 @@ def base_input_view(is_structure): } else: return { - "atoms_types": [["Ga", "As"], ["Ga", "As"]], + "elements": [["Ga", "As"], ["Ga", "As"]], "lattice_vectors": [ 2.8 * (np.ones((3, 3)) - np.eye(3)), 2.9 * (np.ones((3, 3)) - np.eye(3)), ], - "atoms_trajectory": [ + "positions": [ [ [0.0, 0.0, 0.0], [0.25, 0.25, 0.25], @@ -69,20 +83,26 @@ def view(request, not_core): def test_structure_to_view(view: View, Assert, not_core): state = view.to_vasp_viewer().get_state() # check positions, lattice and element types - Assert.allclose(view.positions, state["atoms_trajectory"]) - Assert.allclose(view.elements, state["atoms_types"]) - Assert.allclose(view.lattice_vectors, state["lattice_vectors"]) + assert np.allclose( + view.positions, _trajectory_from_state(state["_atoms_trajectory"]), atol=1e-7 + ) + Assert.allclose(view.elements, state["_atoms_types"]) + assert np.allclose( + view.lattice_vectors, + _trajectory_from_state(state["_lattice_vectors"]), + atol=1e-7, + ) @hasVaspView -@patch("vasp_viewer.Widget", autospec=True) +@patch("vasp.viewer.Widget", autospec=True) def test_ipython(mock_display, view, not_core): display = view._ipython_display_(mode="vasp_viewer") mock_display.assert_called_once() @hasVaspView -@patch("vasp_viewer.Widget", autospec=True) +@patch("vasp.viewer.Widget", autospec=True) def test_ipython_auto(mock_display, view, not_core): display = view._ipython_display_(mode="auto") mock_display.assert_called_once() @@ -93,7 +113,7 @@ def test_ipython_auto(mock_display, view, not_core): def test_camera(view, camera, not_core): view.camera = camera state = view.to_vasp_viewer().get_state() - assert camera == state["selections_camera_mode"] + assert camera == state["_selections_camera_mode"] @hasVaspView @@ -123,8 +143,8 @@ def test_ion_arrows(is_structure, Assert, not_core): ion_arrows=[ IonArrow( np.random.rand( - len(inputs["atoms_trajectory"]), - len(inputs["atoms_trajectory"][0]), + len(inputs["positions"]), + len(inputs["positions"][0]), 3, ), label="Magnetization", @@ -133,8 +153,8 @@ def test_ion_arrows(is_structure, Assert, not_core): ), IonArrow( np.random.rand( - len(inputs["atoms_trajectory"]), - len(inputs["atoms_trajectory"][0]), + len(inputs["positions"]), + len(inputs["positions"][0]), 3, ), label="Velocities", @@ -145,9 +165,13 @@ def test_ion_arrows(is_structure, Assert, not_core): ) state = view.to_vasp_viewer().get_state() for arrow_group_view, arrow_group_state in zip( - view.ion_arrows, state["ion_arrow_groups"] + view.ion_arrows, state["_ion_arrow_groups"] ): - Assert.allclose(arrow_group_view.quantity, arrow_group_state["quantity"]) + assert np.allclose( + arrow_group_view.quantity, + _quantity_from_state(arrow_group_state), + atol=1e-7, + ) assert arrow_group_view.label == arrow_group_state["label"] assert arrow_group_view.color == arrow_group_state["base_color"] assert arrow_group_view.radius == arrow_group_state["base_radius"] @@ -158,8 +182,10 @@ def test_ion_arrows(is_structure, Assert, not_core): def test_supercell(is_structure, Assert, not_core): view = View(**base_input_view(is_structure), supercell=(2, 2, 2)) state = view.to_vasp_viewer().get_state() - Assert.allclose(view.positions, state["atoms_trajectory"]) - Assert.allclose(view.supercell, state["selections_supercell"]) + assert np.allclose( + view.positions, _trajectory_from_state(state["_atoms_trajectory"]), atol=1e-7 + ) + Assert.allclose(view.supercell, state["_selections_bounds"][3:6]) @hasVaspView @@ -170,7 +196,7 @@ def test_supercell(is_structure, Assert, not_core): def test_showcell(is_structure, is_show_cell, not_core): view = View(**base_input_view(is_structure), show_cell=is_show_cell) state = view.to_vasp_viewer().get_state() - assert state["selections_show_lattice"] == is_show_cell + assert state["_selections_show_lattice"] == is_show_cell @hasVaspView @@ -181,10 +207,7 @@ def test_showcell(is_structure, is_show_cell, not_core): def test_showaxes(is_structure, is_show_axes, not_core): view = View(**base_input_view(is_structure), show_axes=is_show_axes) state = view.to_vasp_viewer().get_state() - assert state["selections_show_xyz"] == False - assert state["selections_show_abc"] == is_show_axes - assert state["selections_show_xyz_aside"] == is_show_axes - assert state["selections_show_abc_aside"] == is_show_axes + assert state["_selections_show_abc"] == is_show_axes @hasVaspView @@ -195,8 +218,8 @@ def test_showaxes_different_origin(is_structure, Assert, not_core): **base_input_view(is_structure), show_axes=True, show_axes_at=axes_offset ) state = view.to_vasp_viewer().get_state() - Assert.allclose(state["selections_axes_abc_shift"], axes_offset) - Assert.allclose(state["selections_axes_xyz_shift"], axes_offset) + Assert.allclose(state["_selections_axes_abc_shift"], axes_offset) + Assert.allclose(state["_selections_axes_xyz_shift"], axes_offset) @hasVaspView @@ -204,7 +227,7 @@ def test_showaxes_different_origin(is_structure, Assert, not_core): def test_atom_radius(atom_radius, not_core): view = View(**base_input_view(is_structure=True), atom_radius=atom_radius) state = view.to_vasp_viewer().get_state() - assert state["selections_atom_radius"] == atom_radius + assert state["_selections_atom_radius"] == atom_radius @hasVaspView @@ -213,7 +236,7 @@ def test_structure_title(not_core): **base_input_view(is_structure=True), structure_title="My Structure Title" ) state = view.to_vasp_viewer().get_state() - assert state["selections_descriptor"] == view.structure_title + assert state["_selections_descriptor"] == view.structure_title @hasVaspView diff --git a/tests/third_party/view/test_view.py b/tests/third_party/view/test_view.py index 9d89f96bc..d2bb8e4b0 100644 --- a/tests/third_party/view/test_view.py +++ b/tests/third_party/view/test_view.py @@ -416,3 +416,163 @@ def test_incorrect_shape_raises_error(view): incorrect_unit_cell = np.zeros((len(view.lattice_vectors), 2, 4)) with pytest.raises(exception.IncorrectUsage): View(view.elements, incorrect_unit_cell, view.positions) + + +def test_add_requires_compatible_trajectory_fields(not_core): + left = View(**base_input_view(is_structure=False)) + right = View(**base_input_view(is_structure=False)) + + combined = left + right + + assert np.array_equal(combined.positions, left.positions) + assert np.array_equal(combined.lattice_vectors, left.lattice_vectors) + assert np.array_equal(combined.elements, left.elements) + + with pytest.raises(exception.IncorrectUsage): + left + View(**base_input_view(is_structure=True)) + + +def test_add_merges_scalar_fields_and_raises_on_conflict(not_core): + left = View(structure_title="left title", atom_radius=0.8, **base_input_view(False)) + right = View(structure_title=None, atom_radius=0.8, **base_input_view(False)) + + combined = left + right + + assert combined.structure_title == "left title" + assert combined.atom_radius == 0.8 + + with pytest.raises(exception.IncorrectUsage): + left + View(camera="perspective", **base_input_view(False)) + + +def test_add_combines_grid_scalars_and_ion_arrows(not_core): + number_atoms = len(base_input_view(False)["elements"][0]) + shared_grid = GridQuantity(np.arange(8, dtype=float).reshape(1, 2, 2, 2), "shared") + extra_grid = GridQuantity(np.ones((1, 2, 2, 2)), "extra") + shared_arrow = IonArrow( + quantity=np.ones((1, number_atoms, 3)), + label="forces", + color="#2FB5AB", + radius=0.1, + ) + extra_arrow = IonArrow( + quantity=np.zeros((1, number_atoms, 3)), + label="moments", + color="#4C265F", + radius=0.2, + ) + + left = View( + grid_scalars=[shared_grid], + ion_arrows=[shared_arrow], + **base_input_view(False), + ) + right = View( + grid_scalars=[copy.deepcopy(shared_grid), extra_grid], + ion_arrows=[copy.deepcopy(shared_arrow), extra_arrow], + **base_input_view(False), + ) + + combined = left + right + + assert len(combined.grid_scalars) == 2 + assert len(combined.ion_arrows) == 2 + assert combined.grid_scalars[0].label == "shared" + assert combined.grid_scalars[1].label == "extra" + assert combined.ion_arrows[0].label == "forces" + assert combined.ion_arrows[1].label == "moments" + + +def test_add_combines_special_sequences_as_sequences(not_core): + number_atoms = len(base_input_view(False)["elements"][0]) + grid_left = GridQuantity(np.arange(8, dtype=float).reshape(1, 2, 2, 2), "left") + grid_right = GridQuantity(np.ones((1, 2, 2, 2)), "right") + arrow_left = IonArrow( + quantity=np.ones((1, number_atoms, 3)), + label="left", + color="#2FB5AB", + radius=0.1, + ) + arrow_right = IonArrow( + quantity=np.zeros((1, number_atoms, 3)), + label="right", + color="#4C265F", + radius=0.2, + ) + + left = View( + grid_scalars=(grid_left,), + ion_arrows=(arrow_left,), + **base_input_view(False), + ) + right = View( + grid_scalars=(copy.deepcopy(grid_left), grid_right), + ion_arrows=(copy.deepcopy(arrow_left), arrow_right), + **base_input_view(False), + ) + + combined = left + right + + assert isinstance(combined.grid_scalars, tuple) + assert isinstance(combined.ion_arrows, tuple) + assert tuple(grid.label for grid in combined.grid_scalars) == ("left", "right") + assert tuple(arrow.label for arrow in combined.ion_arrows) == ("left", "right") + + +def test_add_special_sequence_preserves_left_sequence_type(not_core): + number_atoms = len(base_input_view(False)["elements"][0]) + left_grid = [GridQuantity(np.arange(8, dtype=float).reshape(1, 2, 2, 2), "left")] + right_grid = ( + GridQuantity(np.arange(8, dtype=float).reshape(1, 2, 2, 2), "left"), + GridQuantity(np.ones((1, 2, 2, 2)), "right"), + ) + left_arrow = [ + IonArrow( + quantity=np.ones((1, number_atoms, 3)), + label="left", + color="#2FB5AB", + radius=0.1, + ) + ] + right_arrow = ( + IonArrow( + quantity=np.ones((1, number_atoms, 3)), + label="left", + color="#2FB5AB", + radius=0.1, + ), + IonArrow( + quantity=np.zeros((1, number_atoms, 3)), + label="right", + color="#4C265F", + radius=0.2, + ), + ) + + left = View(grid_scalars=left_grid, ion_arrows=left_arrow, **base_input_view(False)) + right = View( + grid_scalars=right_grid, + ion_arrows=right_arrow, + **base_input_view(False), + ) + + combined = left + right + + assert isinstance(combined.grid_scalars, list) + assert isinstance(combined.ion_arrows, list) + assert [grid.label for grid in combined.grid_scalars] == ["left", "right"] + assert [arrow.label for arrow in combined.ion_arrows] == ["left", "right"] + + +def test_add_does_not_trigger_validation(not_core, monkeypatch): + left = View(**base_input_view(False)) + right = View(**base_input_view(False)) + + def _raise_if_called(*_, **__): + raise RuntimeError("validation should not run during View combination") + + monkeypatch.setattr(View, "_verify", _raise_if_called) + + combined = left + right + + assert np.array_equal(combined.positions, left.positions)