Skip to content

Commit

Permalink
Analysis aggregation of multidimensional data (#272)
Browse files Browse the repository at this point in the history
* add aggregation option to Properties1D

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
PythonFZ and pre-commit-ci[bot] authored Nov 9, 2023
1 parent 5740f95 commit 2c13b09
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions zndraw/analyse/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,10 @@ class Properties1D(BaseModel):
smooth: bool = False

model_config = ConfigDict(json_schema_extra=_schema_from_atoms)
aggregation: t.Literal["mean", "median", "max", ""] = Field(
"",
description="For multidimensional data, aggregate over all dimensions, except the first one.",
)

@classmethod
def model_json_schema_from_atoms(cls, schema: dict) -> dict:
Expand All @@ -150,9 +154,19 @@ def model_json_schema_from_atoms(cls, schema: dict) -> dict:
return schema

def run(self, vis):
vis.log("Downloading data...")
atoms_lst = list(vis)
data = np.array([x.calc.results[self.value] for x in atoms_lst])

if data.ndim > 1:
axis = tuple(range(1, data.ndim))
if self.aggregation == "mean":
data = np.mean(data, axis=axis)
elif self.aggregation == "median":
data = np.median(data, axis=axis)
elif self.aggregation == "max":
data = np.max(data, axis=axis)

df = pd.DataFrame({"step": list(range(len(atoms_lst))), self.value: data})

fig = px.line(df, x="step", y=self.value, render_mode="svg")
Expand Down

0 comments on commit 2c13b09

Please sign in to comment.