enable convert_nan API (#690)

* enable `convert_nan` API

* fix lock

Author: PythonFZ

poetry.lock

@@ -23,7 +23,7 @@ celery = "^5"
 sqlalchemy = "^2"
 redis = "^5"
 splines = "^0.3"
-znsocket = "^0.2.6"
+znsocket = "^0.2.7"
 znjson = "^0.2.4"
 
 

pyproject.toml

@@ -98,6 +98,7 @@ class FileIO:
     step: int = 1
     remote: str | None = None
     rev: str | None = None
+    convert_nan: bool = False
 
     def to_dict(self):
         return dataclasses.asdict(self)

zndraw/base.py

@@ -122,6 +122,7 @@ def read_file(fileio: dict) -> None:
         r=current_app.extensions["redis"],
         url=current_app.config["SERVER_URL"],
         token="default",
+        convert_nan=file_io.convert_nan,
     )
     bonds_calculator = ASEComputeBonds()
     if current_app.config.get("COMPUTE_BONDS", False):

zndraw/tasks/__init__.py

@@ -25,7 +25,7 @@ def upload(
     """Upload a file to ZnDraw."""
     if token is None:
         token = str(uuid.uuid4())
-    vis = ZnDraw(url=url, token=token)
+    vis = ZnDraw(url=url, token=token, convert_nan=fileio.convert_nan)
     typer.echo(f"Uploading to: {url}/token/{vis.token}")
 
     if not append:

zndraw/upload.py

@@ -146,6 +146,7 @@ class ZnDraw(ZnDrawBase):
         )
     )
     verify: bool | str = True
+    convert_nan: bool = False
 
     max_atoms_per_call: int = dataclasses.field(default=1000, repr=False)
     # number of `ase.Atom` to send per call
@@ -256,6 +257,7 @@ def __setitem__(
             converter=[ASEConverter],
             socket=self._refresh_client,
             max_commands_per_call=100,
+            convert_nan=self.convert_nan,
         )
         if not self.r.exists(f"room:{self.token}:frames") and self.r.exists(
             "room:default:frames"
@@ -345,6 +347,7 @@ def insert(self, index: int, value: ase.Atoms):
             f"room:{self.token}:frames",
             converter=[ASEConverter],
             socket=self._refresh_client,
+            convert_nan=self.convert_nan,
         )
         if not self.r.exists(f"room:{self.token}:frames") and self.r.exists(
             "room:default:frames"
@@ -392,6 +395,7 @@ def extend(self, values: list[ase.Atoms]):
             converter=[ASEConverter],
             socket=self._refresh_client,
             max_commands_per_call=100,
+            convert_nan=self.convert_nan,
         )
         # TODO: why is there no copy action here?
         show_tbar = (len(values[0]) * len(values)) > self.max_atoms_per_call
@@ -463,6 +467,7 @@ def log(self, message: str) -> None:
             f"room:{self.token}:chat",
             socket=self._refresh_client,
             max_commands_per_call=100,
+            convert_nan=self.convert_nan,
         ).append(msg)
 
     @property

zndraw/zndraw.py

@@ -124,6 +124,11 @@ def main(
     plots: list[str] = typer.Option(
         None, "--plots", "-p", help="List of plots to be shown in the ZnDraw GUI."
     ),
+    convert_nan: bool = typer.Option(
+        False,
+        help="Convert NaN values to None. This is slow and experimental, but if your file contains NaN/inf values, it is required.",
+        envvar="ZNDRAW_CONVERT_NAN",
+    ),
 ):
     """Start the ZnDraw server.
 
@@ -187,6 +192,7 @@ def main(
         start=start,
         stop=stop,
         step=step,
+        convert_nan=convert_nan,
     )
 
     if url is not None: