From fd005a7e50a29f1a7de3deb08cf55435e1b70ca1 Mon Sep 17 00:00:00 2001
From: Guillaume Fraux <guillaume.fraux@epfl.ch>
Date: Thu, 26 Oct 2023 14:57:46 +0200
Subject: [PATCH] Do not put data in ase.Atoms to immediately remove it

---
 python/examples/colors.py | 97 +++++++++++++++------------------------
 1 file changed, 38 insertions(+), 59 deletions(-)

diff --git a/python/examples/colors.py b/python/examples/colors.py
index 56635d272..1e37f2f2a 100644
--- a/python/examples/colors.py
+++ b/python/examples/colors.py
@@ -1,91 +1,70 @@
 """
-Atom property coloring in chemiscope
-====================================
+Atom property coloring
+======================
 
-This example demonstrates how to color atoms based on scalar
-properties.
+This example demonstrates how to color atoms based on scalar properties.
 
-Note that the same parameters can be used with `chemiscope.show`
-to visualize an interactive widget in a Jupyter notebook.
+Note that the same parameters can be used with `chemiscope.show` to visualize an
+interactive widget in a Jupyter notebook.
 """
-
 import ase.io
-import chemiscope
 import numpy as np
 
+import chemiscope
+
 # loads a dataset of structures
 frames = ase.io.read("data/alpha-mu.xyz", ":")
 
-# converts the arrays from the format they are stored in to an array
-# format, and computes scalar quantities to display as atom coloring
-for a in frames:
-    a.arrays["polarizability"] = np.array(
-        [
-            (axx + ayy + azz) / 3
-            for (axx, ayy, azz) in zip(
-                a.arrays["axx"], a.arrays["ayy"], a.arrays["azz"]
-            )
-        ]
-    )
+# compute some scalar quantities to display as atom coloring
+polarizability = []
+alpha_eigenvalues = []
+anisotropy = []
+
+for frame in frames:
+    for axx, ayy, azz, axy, axz, ayz in zip(
+        frame.arrays["axx"],
+        frame.arrays["ayy"],
+        frame.arrays["azz"],
+        frame.arrays["axy"],
+        frame.arrays["axz"],
+        frame.arrays["ayz"],
+    ):
+        polarizability.append((axx + ayy + azz) / 3)
 
-    # one possible measure of anisotropy...
-    a.arrays["alpha_eigenvalues"] = np.array(
-        [
-            np.linalg.eigvalsh([[axx, axy, axz], [axy, ayy, ayz], [axz, ayz, azz]])
-            for (axx, ayy, azz, axy, axz, ayz) in zip(
-                a.arrays["axx"],
-                a.arrays["ayy"],
-                a.arrays["azz"],
-                a.arrays["axy"],
-                a.arrays["axz"],
-                a.arrays["ayz"],
-            )
-        ]
-    )
+        # one possible measure of anisotropy...
+        eigenvalues = np.linalg.eigvalsh(
+            [[axx, axy, axz], [axy, ayy, ayz], [axz, ayz, azz]]
+        )
+        alpha_eigenvalues.append(eigenvalues)
 
-    a.arrays["anisotropy"] = (
-        a.arrays["alpha_eigenvalues"][:, 2] - a.arrays["alpha_eigenvalues"][:, 0]
-    )
+        anisotropy.append(eigenvalues[2] - eigenvalues[0])
 
 
-# now we just write the chemiscope datafile
+# now we just write the chemiscope input file
 chemiscope.write_input(
     "colors-example.json.gz",
     frames=frames,
     # properties can be extracted from the ASE.Atoms frames
-    properties=chemiscope.extract_properties(
-        frames, only=["polarizability", "anisotropy", "alpha_eigenvalues"]
-    ),
+    properties={
+        "polarizability": np.vstack(polarizability),
+        "anisotropy": np.vstack(anisotropy),
+        "alpha_eigenvalues": np.vstack(alpha_eigenvalues),
+    },
     # the write_input function also allows defining the default visualization settings
-    settings={  
+    settings={
         "map": {
             "x": {"property": "alpha_eigenvalues[1]"},
             "y": {"property": "alpha_eigenvalues[2]"},
             "z": {"property": "alpha_eigenvalues[3]"},
-            "palette": "seismic",
             "color": {"property": "anisotropy"},
         },
         "structure": [
             {
-                "spaceFilling": False,
-                "atomLabels": False,
-                "atoms": True,
-                "axes": "off",
-                "keepOrientation": False,
-                "playbackDelay": 700,
-                "environments": {
-                    "activated": True,
-                    "bgColor": "CPK",
-                    "bgStyle": "licorice",
-                    "center": False,
-                    "cutoff": 4,
-                },
-                # these are the color visualization options
-                "color": {"property": "anisotropy", "min": 1, "max": 15},
+                "color": {"property": "anisotropy"},
             }
         ],
     },
-    # these are atomic properties - in order to view them in the map panel, we must 
-    # indicate that all atoms are environment centers
+    # the properties we want to visualise are atomic properties - in order to view them
+    # in the map panel, we must indicate that all atoms are environment centers
     environments=chemiscope.all_atomic_environments(frames),
 )