import numpy as np
from .mesh_data import MeshData
from .plane_data import PlaneData
###############################################################################
# The "Accessor_Mixin" class is a mixin that provides methods to access
# any variable in 2d, 3d, f2d, and f3d data structures.
###############################################################################
[docs]
class ArrayAccessorMixin:
"""
Mixin class to provide access to data arrays in various formats.
This class should be used with classes that have a `data` attribute
structured as a dictionary of dictionaries, where the outer key is the variable name
and the inner key is the step index.
Design note
-----------
This mixin provides generic access helpers that only assume the common
``self.data[var_name][step_index]`` storage layout. Reader classes should
add thin semantic wrappers (for example ``get_mesh_data`` or
``get_distribution``) on top of these methods.
"""
[docs]
def has_var(self, var_name):
"""Return ``True`` if ``var_name`` exists in ``self.data``."""
return var_name in self.data
[docs]
def list_vars(self):
"""Return sorted variable names stored in ``self.data``."""
return sorted(self.data.keys())
[docs]
def list_step_indices(self, var_name=None):
"""
Return sorted step/file-index keys.
Parameters
----------
var_name : str or None
If provided, return keys only for that variable. If ``None``, return
the sorted union of keys across all variables.
"""
if var_name is not None:
if var_name not in self.data:
raise KeyError(f"Variable '{var_name}' not found in data.")
return sorted(self.data[var_name].keys())
keys = set()
for step_dict in self.data.values():
keys.update(step_dict.keys())
return sorted(keys)
[docs]
def get_item(self, var_name, step_index):
"""Return the raw stored object for ``var_name`` at ``step_index``."""
if var_name not in self.data:
raise KeyError(f"Variable '{var_name}' not found in data.")
if step_index not in self.data[var_name]:
raise KeyError(
f"Step/file index '{step_index}' not found for variable '{var_name}'."
)
return self.data[var_name][step_index]
[docs]
def get_as(self, var_name, step_index, expected_type):
"""
Return a stored item and validate its type.
Parameters
----------
expected_type : type | tuple[type, ...]
Allowed Python class/type(s).
"""
obj = self.get_item(var_name, step_index)
if not isinstance(obj, expected_type):
if isinstance(expected_type, tuple):
type_name = ", ".join(t.__name__ for t in expected_type)
else:
type_name = expected_type.__name__
raise TypeError(
f"Variable '{var_name}' at step/file index {step_index} has type "
f"{type(obj).__name__}; expected {type_name}."
)
return obj
[docs]
def get_array(self, var_name):
"""
Returns the raw NumPy array for all steps of a given variable.
Args:
var_name (str):
Name of the variable to extract (e.g., "eden", "time", etc.).
steps: (list, optional):
List of specific steps to include in the output.
If None, all available steps will be used.
Returns:n
np.ndarray: The raw NumPy array for the specified variable across all steps.
- For scalar variables (e.g., time), the shape will be (n_steps,).
- For mesh-based variables, the shape will be (n_steps, n_planes, n_vertices).
- For plane-based variables, the shape will be (n_steps, n_nodes). (Or something similar.)
"""
if var_name not in self.data:
raise KeyError(f"Variable '{var_name}' not found in data.")
step_dict = self.data[var_name]
step_keys = sorted(step_dict.keys()) #f steps is None else sorted(steps)
first_val = step_dict[step_keys[0]]
# Case 1: scalar or time
if var_name == "time" or not hasattr(first_val, "get_data"):
return np.array([step_dict[k] for k in step_keys])
# Case 2: MeshData or PlaneData with .get_data()
return np.stack([step_dict[k].get_data() for k in step_keys])