diff --git a/src/osyris/__init__.py b/src/osyris/__init__.py index 8191f7f..06167dd 100644 --- a/src/osyris/__init__.py +++ b/src/osyris/__init__.py @@ -7,3 +7,4 @@ from .units import units from .core import Array, Datagroup, Dataset, Plot, Vector from .plot import histogram1d, histogram2d, scatter, map, plot +from .spatial import extract_box, extract_sphere diff --git a/src/osyris/spatial/__init__.py b/src/osyris/spatial/__init__.py new file mode 100644 index 0000000..4f59a55 --- /dev/null +++ b/src/osyris/spatial/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2022 Osyris contributors (https://github.com/osyris-project/osyris) + +# flake8: noqa + +from .subdomain import extract_sphere, extract_box diff --git a/src/osyris/spatial/subdomain.py b/src/osyris/spatial/subdomain.py new file mode 100644 index 0000000..3ea9003 --- /dev/null +++ b/src/osyris/spatial/subdomain.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2022 Osyris contributors (https://github.com/osyris-project/osyris) + +import numpy as np +from .. import Dataset +import warnings + + +def extract_sphere(dataset, radius, origin): + """ + Extract a spherical subdomain around an origin point. + """ + subdomain = Dataset() + subdomain.meta = dataset.meta.copy() + + for name, group in dataset.items(): + pos = group.get("position", group.parent["amr"]["position"]) + if pos.shape != group.shape: + warnings.warn( + "Ignoring datagroup '{}', which has no position ".format(group) + + "vector and has different shape than 'amr' group.") + continue + r = (pos - origin).norm + c = (r < radius).values + if np.any(c): + subdomain[name] = dataset[name][c] + + return subdomain + + +def extract_box(dataset, dx, dy, dz, origin): + """ + Extract a cubic domain of size dx, dy & dz around an origin point + """ + subdomain = Dataset() + subdomain.meta = dataset.meta.copy() + + for name, group in dataset.items(): + pos = group.get("position", group.parent["amr"]["position"]) + if pos.shape != group.shape: + warnings.warn( + "Ignoring datagroup '{}', which has no position ".format(group) + + "vector and has different shape than 'amr' group.") + continue + centered_pos = pos - origin + cx = (centered_pos.x <= dx * .5) & (centered_pos.x >= -dx * .5) + cy = (centered_pos.y <= dy * .5) & (centered_pos.y >= -dy * .5) + cz = (centered_pos.z <= dz * .5) & (centered_pos.z >= -dz * .5) + c = (cx & cy & cz).values + if np.any(c): + subdomain[name] = dataset[name][c] + + return subdomain