Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
243 commits
Select commit Hold shift + click to select a range
39488ab
pscadios2: add original, not split into components, data
germasch Dec 21, 2024
4f627e8
test_adios2: add test_open_twice()
germasch Dec 21, 2024
f888880
test_xarray_adios2: test some selections
germasch Dec 21, 2024
9b57c4b
test_xarray_adios2: add test_computed{,_via_lambda}
germasch Dec 21, 2024
1068a2f
add some logging
germasch Dec 21, 2024
540f30e
pscadios2: no need to import anymore
germasch Dec 21, 2024
0fff803
tests: add pfd_moments test data
germasch Dec 21, 2024
aad6a6b
tests: test opening pfd_moments
germasch Dec 21, 2024
b75c6ca
pscadios2: add per-component variables later
germasch Dec 21, 2024
a265a6e
import psc, rather than symbols from it
germasch Dec 21, 2024
a4527e1
import adios2py, rather than symbols from it
germasch Dec 21, 2024
866414e
pscadios2: add coords later
germasch Dec 21, 2024
4a2f1f7
pscadios2: don't need species_names, etc in most places
germasch Dec 21, 2024
09a42fa
PscAdios2Array: don't need orig_varname, component anymore
germasch Dec 21, 2024
bf0994f
pscadios2: remove some of the psc-prefixed naming
germasch Dec 21, 2024
2e67e2d
do psc_open_dataset inline
germasch Dec 21, 2024
f44aa38
_normalize_path first
germasch Dec 21, 2024
31a7ac0
follow standard way of open_dataset()
germasch Dec 21, 2024
c323558
adios2py: special case reading of string attributes
germasch Dec 22, 2024
b1217e6
pscadios2: keep variable-associated attributes with the variable
germasch Dec 22, 2024
54efd02
pscadios2: set dims based on xr-dims attribute if possible
germasch Dec 22, 2024
3fe1339
adios2py: since we're keeping track of opened variables, use as cache
germasch Dec 22, 2024
fbe596c
adios2py: add "is_reverse_dims"
germasch Dec 22, 2024
df1bb53
pscadios2: reverse dims if written in Fortran
germasch Dec 22, 2024
13d4a09
pscadios2: rename psc field component to comp_{var_name}
germasch Dec 22, 2024
4932f3d
pscadios2: make the tests independent of actual ordering
germasch Dec 22, 2024
495b1de
pscadios2: make psc component selection ordering-independent
germasch Dec 22, 2024
66e7fe4
adios2py: implement conditional reversing dims
germasch Dec 22, 2024
9444aba
adios2py: reverse_dims based on IsReverseDims flag
germasch Dec 22, 2024
f7ad096
minor cleanup
germasch Dec 22, 2024
77dcfc7
adios2py: determine reverse_dims override once
germasch Dec 22, 2024
56f983b
adios2py: allow for creating File from pre-existing io/engine
germasch Dec 22, 2024
2454f01
pascadios2: allow passing of already opened DataStore
germasch Dec 22, 2024
105bd52
pscadios2: dims are now written consistently in native order
germasch Dec 22, 2024
79de2a1
adios2py: support reading single value
germasch Dec 22, 2024
1d52e07
add support for selecting / keeping track of step
germasch Dec 24, 2024
6592ae4
adios2py: raise exception if variable not found
germasch Dec 24, 2024
9d77499
adios2py: add .num_steps(), .begin_step(), .end_step() and tests
germasch Dec 24, 2024
2d259c9
adios2py add test_variable_closed()
germasch Dec 24, 2024
a3bc088
adios2py: use "pfd_file" fixture to open sample file
germasch Dec 24, 2024
94fc104
adios2py: add Variable.__bool__
germasch Dec 24, 2024
b83a410
adios2py: make shape() a property
germasch Dec 24, 2024
1522124
adios2py: make name and dtype properties
germasch Dec 24, 2024
0dd361f
adios2py: add test for repr()
germasch Dec 24, 2024
777faef
adios2py: standardize Variable.__repr__ a bit
germasch Dec 24, 2024
2367683
adios2py: Variable.__repr__ don't raise if called on closed variable
germasch Dec 24, 2024
99e7f1d
adios2py: add File.__repr__
germasch Dec 24, 2024
72c7d8d
adios2py: add tests for getitem_{scalar,arr1d}
germasch Dec 24, 2024
779feeb
adios2py: implement Variable.__array__
germasch Dec 24, 2024
8a4b8fa
adios2py test: provide File.steps() iterator
germasch Dec 24, 2024
20a5be4
adios2py/test: use test_file fixture for read_streaming_adios2py
germasch Dec 24, 2024
c84501c
adios2py: don't make assumptions in steps() about starting from the b…
germasch Dec 24, 2024
60e1e46
adios2py: yield self from steps()
germasch Dec 24, 2024
e3d19cd
use new step iterator
germasch Dec 24, 2024
8a5a4c9
adios2py: add tests for step persistence (failing)
germasch Dec 24, 2024
2a4f56e
adios2py: add test for current_step() outside of begin/end step (xfail)
germasch Dec 25, 2024
992c8ae
adios2py: break out _update_variables_attributes()
germasch Dec 25, 2024
1c3e3df
adios2py: introduce reset() and clear _open_vars on begin_step()
germasch Dec 25, 2024
cc8e9ca
Adios2Store: override load()
germasch Dec 25, 2024
caed549
adios2store: start some tests
germasch Dec 25, 2024
aa90fa2
adios2py: add optional io= arg to adios2py.File()
germasch Dec 25, 2024
0182105
adios2py: support optional "parameters" arg
germasch Dec 25, 2024
ea9280c
Adios2Store: optionaly pass parameters
germasch Dec 25, 2024
c21ed19
adios2py/Adios2Store: support "engine" keyword
germasch Dec 25, 2024
8b0af64
adios2py: add _engine and _io properties
germasch Dec 25, 2024
87d2687
adios2py: implement File.__bool__
germasch Dec 25, 2024
f9fef17
adios2py: remove io argument to File() again
germasch Dec 25, 2024
446679d
adios2py: separate out generator for making unique IO objects
germasch Dec 25, 2024
5d9e3d5
adios2py: make _state non optional
germasch Dec 25, 2024
3e964ac
adios2py: get rid of FileState.is_open
germasch Dec 25, 2024
e4f6b56
adios2py: move FileState.close() -> File.close()
germasch Dec 25, 2024
9c786d6
adios2py: keep only data in FileState
germasch Dec 25, 2024
b1363f5
adios2py.File: rename ._io -> .io, ._engine -> .engine
germasch Dec 25, 2024
27f00c9
adios2py: mv .io_name -> File._io_name
germasch Dec 25, 2024
9600696
adios2py: mv .io -> File._io
germasch Dec 25, 2024
d3917a6
adios2py: FileState is now gone
germasch Dec 25, 2024
56a5089
adios2py: mark _io and _engine as optional
germasch Dec 25, 2024
3a2f832
adios2py: explicitly use flag for io/engine ownership
germasch Dec 25, 2024
b708b5c
adiospy: _io_generator doesn't need to yield io_name
germasch Dec 25, 2024
130844b
adios2py.File: add defaults for member vars
germasch Dec 25, 2024
7b947c3
adios2py/Adios2Store: rename engine arg -> engine_type
germasch Dec 25, 2024
6b0af67
adios2py.File: .variable_names -> .keys()
germasch Dec 25, 2024
f101de5
adios2py: don't cache _variable_names
germasch Dec 25, 2024
6d005b3
adios2py: use .attrs.keys() instead of .attribute_names
germasch Dec 25, 2024
6c144da
adios2py: don't cache attribute_names
germasch Dec 25, 2024
f147586
adios2py: keep reference to file rather than engine
germasch Dec 25, 2024
87753f7
adios2py: create adios2_var in Variable()
germasch Dec 25, 2024
ea7c431
adios2py.Variable: access through .var for the most part
germasch Dec 25, 2024
ee1a50a
adios2py.Variable: don't cache adios2.Variable
germasch Dec 25, 2024
e15a213
adios2py.Variable: inline _set_selection(), get var as local
germasch Dec 25, 2024
eec8c7b
adios2py.Variable: don't expose is_reverse_dims
germasch Dec 25, 2024
ec61d85
adios2py.Variable: some cleanup
germasch Dec 25, 2024
4a7b7ab
adios2py.Variable: .step -> ._step and get rid of all the open_vars l…
germasch Dec 25, 2024
e39bb64
adios2py: remove reverse_dims override for PSC
germasch Dec 25, 2024
c54302e
adios2py.File: access attributes through .attrs
germasch Dec 25, 2024
273a90b
adios2py.File: add attrs.__contains__
germasch Dec 25, 2024
58d92f3
adios2py.File.attrs: implement Mapping
germasch Dec 25, 2024
b0ddfae
psc: no more _get_array_attribute()
germasch Dec 25, 2024
e3027bb
adios2py: cosmetic
germasch Dec 25, 2024
7f35188
tests/adios2py: separate out test_filename fixture
germasch Dec 25, 2024
b7ae46e
adios2py: start random access opening
germasch Dec 25, 2024
0d1b67e
adios2py: toward supporting random access read
germasch Dec 25, 2024
3ef8450
adios2py: add test for current rra access hack
germasch Dec 25, 2024
007cfff
adios2py: use step._step to get currently iterated step
germasch Dec 25, 2024
fe6159d
adios2py: iterating rra steps works, if hacky...
germasch Dec 25, 2024
4a06d91
adios2py: __getitem__ instead of get_variable()
germasch Dec 26, 2024
d0e9c33
adios2py: raise KeyError on variable not found
germasch Dec 26, 2024
5480924
adios2py.File: inherit from mapping
germasch Dec 26, 2024
70d35d2
adios2py: simplify two uses of .keys()
germasch Dec 26, 2024
fbc6b32
pscadios2: support mode "rra" in open_existing()
germasch Dec 26, 2024
bdbbd4d
adios2py/Adios2Store: don't need open_existing anymore
germasch Dec 26, 2024
0f543b6
adios2py: get rid of .reset()
germasch Dec 26, 2024
736c937
adios2py: test_read_streaming_adios2py_step_persist() with explicit mode
germasch Dec 26, 2024
57959f1
adios2py/tests: rra mode delayed var works, r mode does not
germasch Dec 26, 2024
fa8e33f
adios2py: implement StepsProxy
germasch Dec 26, 2024
78c91ea
adios2py: implement File.steps.__len__
germasch Dec 26, 2024
da57600
adios2py: File.num_steps() -> len(File.steps)
germasch Dec 26, 2024
cd36046
adios2py: add File.steps[n] to select nth step
germasch Dec 26, 2024
6d72f29
adios2py: raise when trying to select step outside random access mode
germasch Dec 26, 2024
1ea54a7
adios2py: _own_io_engine unused -> gone
germasch Dec 26, 2024
d8cd62f
adios2py: and now we're back to a separate FileState class...
germasch Dec 26, 2024
3b368ca
adios2py: open io and engine in FileState.__init__
germasch Dec 26, 2024
e05d65f
adios2py: filename and mode -> FileState
germasch Dec 26, 2024
2ea24c5
adios2py: Introduce Group < Step, File class hiearchy
germasch Dec 26, 2024
baf8886
pscadios2: add Adios2Store.set_step()
germasch Dec 26, 2024
6d3c401
adios2py: add Step.current_step()
germasch Dec 26, 2024
906124b
tests/Adios2Store: split up fixture
germasch Dec 26, 2024
0933b5b
tests/Adios2Store: test random access with .set_step()
germasch Dec 26, 2024
0f4ad82
Adios2Store: Slightly less ugly way of selecting step
germasch Dec 26, 2024
f7a2db9
Adios2Store: ._step -> .step for consistency
germasch Dec 26, 2024
a88ea46
adios2py: some test cleanup
germasch Dec 26, 2024
fdb9d7c
adios2py: File doesn't need set_step(), __getitem__()
germasch Dec 26, 2024
ea42ff2
adios2py: add begin_step(), end_step() to FileState
germasch Dec 26, 2024
2e74b59
adios2py: remove File.{begin,end}_step()
germasch Dec 26, 2024
b9ef697
adios2py: File.steps.next() returning context manager
germasch Dec 26, 2024
7e95d89
adios2py: keep track of step_status; fancy generator / ctx manager stuff
germasch Dec 26, 2024
3580665
adios2py: use actual __next__
germasch Dec 26, 2024
a5b8edc
FileState: keep track of current_step
germasch Dec 26, 2024
2ec151d
adios2py: Step() now always as an actual step (not None)
germasch Dec 26, 2024
550e050
adios2py: keep reference to FileState in StepsProxy
germasch Dec 26, 2024
4dcb383
adios2py: no more File.current_step()
germasch Dec 26, 2024
0c807c9
adios2py: always set step in Variable when gotten from Step
germasch Dec 26, 2024
5de7975
adios2py: change exception to KeyError for closed var
germasch Dec 26, 2024
240c493
adios2py.Variable: keep reference FileState
germasch Dec 26, 2024
e25d69f
adios2py: another xfail fixed
germasch Dec 26, 2024
6065e52
adios2py: one more xfail fixed
germasch Dec 26, 2024
5084464
adios2py: use ValueError for closed FileState
germasch Dec 26, 2024
a5497bf
adios2py: no more Group.{io,engine}
germasch Dec 26, 2024
3617457
adios2py: leave checking for closed file to FileState
germasch Dec 26, 2024
898c0bb
adios2py: no more need for the @property .state guard
germasch Dec 26, 2024
c94f76d
adios2py: pass FileState to AttrsProxy
germasch Dec 26, 2024
10b7751
adios2py: __repr__ cleanup
germasch Dec 26, 2024
6aff6b8
tests: fix typo
germasch Dec 26, 2024
3a8b9bf
tests: open_dataset in r / rra modes
germasch Dec 27, 2024
9c5825c
adios2py: simplify Group.__getitem__()
germasch Dec 27, 2024
abcdfc8
adios2py: be more strict about acessing step-free vars inside of a step
germasch Dec 27, 2024
78f50e5
tests/xarray_adios2: merge two tests
germasch Dec 27, 2024
a0c8398
adios2py: move .attrs -> Group
germasch Dec 27, 2024
9db3a9e
pscadios2: allow making Adios2Store from adios2py.{File,Step}
germasch Dec 27, 2024
dacf24c
Adios2Store: stop supporting full File -- for now, Step only
germasch Dec 27, 2024
1684957
xarray_adios2: allow passing adios2py.Group() directly to open_dataset
germasch Dec 27, 2024
ea972ee
adios2py: recognize adios2py.Group in guess_can_open()
germasch Dec 27, 2024
26e4b16
adios2py: call close on FileState::__del__
germasch Dec 27, 2024
d1f54f0
tests: quiet some prints
germasch Dec 27, 2024
e604d28
pscadios2: open_dataset() use rra mode
germasch Dec 27, 2024
d412d87
tests/adios2py: use rra mode to open pfd file
germasch Dec 27, 2024
e080a85
tests/adios2py: do some testing on mixing writing outside/inside of step
germasch Dec 27, 2024
ca4ee16
adios2py: File[var] only in rra mode
germasch Dec 27, 2024
e71e0b9
tests/adios2: fix up file naming
germasch Dec 27, 2024
fed6383
adios2py.Variable: split up __getitem__, _getitem_step_selection
germasch Dec 28, 2024
f93c874
adios2py: mv all of step_selection logic
germasch Dec 28, 2024
4d26f78
adios2py: toward reading all steps in rra mode
germasch Dec 28, 2024
2f7094a
adios2py: very hacky, but seemingly working
germasch Dec 28, 2024
c3585f9
adios2py: we should always get slices or ints, anyway
germasch Dec 28, 2024
34bf540
adios2py: rearrange a bit, add assert
germasch Dec 28, 2024
a04b5f0
adios2py: always pass step selection
germasch Dec 28, 2024
373834b
adios2py: pass step_slice rather than step_selection
germasch Dec 28, 2024
d6ae312
adios2py: pass slice as first of args
germasch Dec 28, 2024
41ac28c
adios2py: renaming
germasch Dec 28, 2024
238a190
adios2py: make step slicing similar to spatial slicing
germasch Dec 28, 2024
c15e786
adios2py: sel_start and sel_stop as lists
germasch Dec 28, 2024
526c520
adios2py: more assimilation
germasch Dec 28, 2024
623a450
adios2py: more assimilation
germasch Dec 28, 2024
8aec605
adios2py: step and spatial selection unified
germasch Dec 28, 2024
1a63d86
adios2py: some fixes and cleanup
germasch Dec 28, 2024
53055ee
adios2py: simpler var_shape
germasch Dec 28, 2024
04d90b4
adios2py: combine all into one loop
germasch Dec 28, 2024
b92567a
adios2py: yet more cleanup
germasch Dec 28, 2024
8f3b2a6
tests/adios2py: check "arr1d" read for all steps
germasch Dec 28, 2024
346d654
tests/adios2py: some more testing -- looks good
germasch Dec 28, 2024
62cf4fe
tests/adios2py: slightly prettify test
germasch Dec 28, 2024
79d8a05
adios2py: streaming mode complains after all on set_step_selection()
germasch Dec 28, 2024
59488ee
adios2py: don't special case rra read with only a single step
germasch Dec 28, 2024
65ea5b0
adios2py: step handling
germasch Dec 28, 2024
77636e4
pscadios2: automatically squeeze "redundant" dim (FIXME?)
germasch Dec 28, 2024
ade1d99
pscadios2: When reading step, remove 0th dimensions name
germasch Dec 28, 2024
1185b0d
pscadios2: add x coord to test, fix redundant handling
germasch Dec 28, 2024
a0f74b7
adios2py: simplify and fix reversing
germasch Dec 29, 2024
01d407d
separate decode_psc function
germasch Dec 31, 2024
5b6be49
pscadios2: start decode_openggcm function
germasch Dec 31, 2024
4a9dbd5
decode_psc fixup
germasch Dec 31, 2024
5f47ef3
pscadios2: some typing work
germasch Dec 31, 2024
ba52d06
pscadios2: do conversion of 7-int time to np.datetime64 in decode_ope…
germasch Dec 31, 2024
a56a41e
pscadios2: handle case of scalar 7-int "time" as well
germasch Dec 31, 2024
f331c6f
tests: add test for "seconds since" decoding
germasch Jan 1, 2025
8702b4d
pscadios2: prepare for updated "dimensions" handling
germasch Jan 2, 2025
6f46d62
pscadios2: drop extra dim if necessary
germasch Jan 2, 2025
9faba19
pscadios2: add _encode_openggcm() for time_array encoding
germasch Jan 3, 2025
fe6ecc0
time_array: add "units" / use for encoding, too
germasch Jan 3, 2025
62684b6
pscadios2: encode_openggcm_variable()
germasch Jan 3, 2025
96b581e
pscadios2: decode openggcm when still at vars, attrs
germasch Jan 3, 2025
b84cf26
pscadios2: use new time_array decoding
germasch Jan 3, 2025
28b9581
pscadios2: renaming
germasch Jan 3, 2025
6d958ac
pscadios2: some conversion improvements
germasch Jan 3, 2025
52e916c
pscadios2: hackily support writing through Adios2Store
germasch Jan 4, 2025
c57d0ff
pscadios2: fix CF encoding in store
germasch Jan 4, 2025
11dba9a
align with stand-alone adios2py
germasch Jan 8, 2025
f3e07b8
pscadios2: stop supporting open_dataset(Adios2Group)
germasch Jan 9, 2025
d76e8c9
test_Adios2Store: avoid ._state
germasch Jan 9, 2025
2686fc3
adios2py: add File.filename property
germasch Jan 9, 2025
e0b70e3
test_adios2py: skip reverse_dims test
germasch Jan 9, 2025
b2fbfd6
pscadios2: use per-var .attrs if available
germasch Jan 9, 2025
ac81d5f
adios2py: switch to using external adios2py
germasch Jan 9, 2025
dd15237
adios2py: remove included adios2py
germasch Jan 9, 2025
79b7ec0
adios2py: only need to support new per-var attributes
germasch Jan 9, 2025
d813515
pscadios2: do begin/end_step from outside of store
germasch Jan 9, 2025
63af90e
test_xarray_adios2/Adios2Store: use adios2py for writing
germasch Jan 9, 2025
ead5182
Adios2Store: add test_dump_to_store(), implement
germasch Jan 10, 2025
6cb0d4b
pscadios2: use new adios2py API
germasch Jan 10, 2025
f12c528
pscadios2/store: use filename_or_obj rather than separate writer
germasch Jan 10, 2025
fe6d401
assimilate toward xarray-adios2
germasch Jan 11, 2025
65805f4
make more tests work with xarray-adios2
germasch Jan 11, 2025
aec8d12
all of the tests except psc-specific work with xarray-adios2
germasch Jan 11, 2025
1b02bf2
update min to python 3.10
germasch Jan 11, 2025
19bae65
RunInfo based on already-read ds, not adios2py File/Step
germasch Jan 11, 2025
f15d54c
psc: make decode_psc work after xarray-adios2 reading
germasch Jan 11, 2025
36bcb0d
psc: mv decode_psc out of pscadios2
germasch Jan 11, 2025
4e8b035
pscpy: don't provide pscadios2_engine
germasch Jan 12, 2025
033ea38
pscadios2: remove pscadios2 altogether
germasch Jan 12, 2025
5b978e5
clean: remove _cori.py
germasch Jan 14, 2025
8e2e087
fixup adios2store gone
germasch Jan 14, 2025
d3059f0
devops: add "xarray-adios2" to dependencies
germasch Jan 14, 2025
01af5e4
tests: remove test-only dependency on ggcmpy
germasch Jan 14, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ repos:
additional_dependencies:
- pytest
- xarray
- adios2
- adios2py

- repo: https://github.com/codespell-project/codespell
rev: "v2.3.0"
Expand Down
12 changes: 6 additions & 6 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ authors = [
description = "Python utilities for reading PSC data"
readme = "README.md"
license.file = "LICENSE"
requires-python = ">=3.8"
requires-python = ">=3.10"
classifiers = [
"Development Status :: 1 - Planning",
"Intended Audience :: Science/Research",
Expand All @@ -23,7 +23,7 @@ classifiers = [
"Typing :: Typed",
]
dynamic = ["version"]
dependencies = ["xarray", "adios2", "typing-extensions"]
dependencies = ["xarray", "xarray-adios2", "typing-extensions"]

[project.optional-dependencies]
test = ["pytest >=6", "pytest-cov >=3"]
Expand All @@ -42,9 +42,6 @@ Homepage = "https://github.com/psc-code/pscpy"
Discussions = "https://github.com/psc-code/pscpy/discussions"
Changelog = "https://github.com/psc-code/pscpy/releases"

[project.entry-points."xarray.backends"]
pscadios2_engine = "pscpy.pscadios2:PscAdios2BackendEntrypoint"

[build-system]
build-backend = "setuptools.build_meta"
requires = ["setuptools>=42", "setuptools-scm>=7"]
Expand Down Expand Up @@ -81,7 +78,7 @@ report.exclude_also = ['\.\.\.', 'if typing.TYPE_CHECKING:']

[tool.mypy]
files = ["src", "tests"]
python_version = "3.8"
python_version = "3.10"
warn_unused_configs = true
strict = true
enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
Expand Down Expand Up @@ -147,3 +144,6 @@ messages_control.disable = [
"missing-function-docstring",
"wrong-import-position",
]

[tool.uv]
reinstall-package = ["pscpy"]
9 changes: 6 additions & 3 deletions src/pscpy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,14 @@

import pathlib

from pscpy import pscadios2 # noqa: F401

from ._version import version as __version__
from .psc import decode_psc

sample_dir = pathlib.Path(__file__).parent / "sample"


__all__ = ["__version__"]
__all__ = [
"__version__",
"decode_psc",
"sample_dir",
]
4 changes: 0 additions & 4 deletions src/pscpy/_core.pyi

This file was deleted.

196 changes: 0 additions & 196 deletions src/pscpy/adios2py/__init__.py

This file was deleted.

73 changes: 51 additions & 22 deletions src/pscpy/psc.py
Original file line number Diff line number Diff line change
@@ -1,23 +1,12 @@
from __future__ import annotations

from typing import Any, Iterable
from collections.abc import Iterable
from typing import Any

import numpy as np
import xarray as xr
from numpy.typing import ArrayLike, NDArray

from .adios2py import File


def _get_array_attribute(
file: File, attribute_name: str, default: ArrayLike | None
) -> NDArray[np.floating[Any]]:
if attribute_name in file.attribute_names:
return np.asarray(file.get_attribute(attribute_name))
if default is not None:
return np.asarray(default)
error_messsage = f"Missing attribute '{attribute_name}' with no default specified."
raise KeyError(error_messsage)


class RunInfo:
"""Global information about the PSC run
Expand All @@ -28,23 +17,22 @@ class RunInfo:

def __init__(
self,
file: File,
ds: xr.Dataset,
length: ArrayLike | None = None,
corner: ArrayLike | None = None,
) -> None:
assert len(file.variable_names) > 0
var = next(iter(file.variable_names))
self.gdims = np.asarray(file.get_variable(var).shape)[0:3]
first_var = ds[next(iter(ds))]
self.gdims = np.asarray(first_var.shape)[::-1][:3]

self.length = _get_array_attribute(file, "length", length)
self.corner = _get_array_attribute(file, "corner", corner)
self.length = ds.attrs.get("length", length)
self.corner = ds.attrs.get("corner", corner)

self.x = self._get_coord(0)
self.y = self._get_coord(1)
self.z = self._get_coord(2)

def _get_coord(self, coord_idx: int) -> NDArray[np.floating[Any]]:
return np.linspace(
def _get_coord(self, coord_idx: int) -> NDArray[Any]:
return np.linspace( # type: ignore[no-any-return]
start=self.corner[coord_idx],
stop=self.corner[coord_idx] + self.length[coord_idx],
num=self.gdims[coord_idx],
Expand Down Expand Up @@ -101,3 +89,44 @@ def get_field_to_component(species_names: Iterable[str]) -> dict[str, dict[str,
)

return field_to_component


def decode_psc(
ds: xr.Dataset,
species_names: Iterable[str],
length: ArrayLike | None = None,
corner: ArrayLike | None = None,
) -> xr.Dataset:
da = ds[next(iter(ds))] # first dataset
if da.dims[0] == "dim_0_1":
# for compatibility, if dimensions weren't saved as attribute in the .bp file,
# fix them up here
ds = ds.rename_dims(
{
da.dims[0]: "step",
da.dims[1]: f"comp_{da.name}",
da.dims[2]: "z",
da.dims[3]: "y",
da.dims[4]: "x",
}
)
ds = ds.squeeze("step")
field_to_component = get_field_to_component(species_names)

data_vars = {}
for var_name in ds:
if var_name in field_to_component:
for field, component in field_to_component[var_name].items(): # type: ignore[index]
data_vars[field] = ds[var_name].isel({f"comp_{var_name}": component})
ds = ds.assign(data_vars)

if length is not None:
run_info = RunInfo(ds, length=length, corner=corner)
coords = {
"x": ("x", run_info.x),
"y": ("y", run_info.y),
"z": ("z", run_info.z),
}
ds = ds.assign_coords(coords)

return ds
Loading
Loading