Fix error loading custom models

ovito-org · Aug 18, 2023 · 28d5952 · 28d5952
1 parent 17c6819
commit 28d5952
Showing 1 changed file with 13 additions and 6 deletions.
diff --git a/src/scoreBasedDenoising/__init__.py b/src/scoreBasedDenoising/__init__.py
@@ -186,7 +186,7 @@ def _modify(self, data, frame, **kwargs):
             case "FCC" | "BCC" | "HCP":
                 model = self.setupFccBccHcpModel(data)
             case "Custom":
-                model = self.setupCustomModel(data)
+                model = self.setupCustomModel()
             case _:
                 raise NotImplementedError
 
@@ -195,12 +195,19 @@ def _modify(self, data, frame, **kwargs):
 
         noisy_atoms = ovito_to_ase(data)
 
-        if self.scale is not None:
-            modelScale = ScoreBasedDenoising.originalScale[self.structure] / self.scale
+        if self.structure == "Custom":
+            modelScale = self.scale
         else:
-            estNNdist = self.estimateNearestNeighborsDistance(data)
-            print(f"Estimated nearest neighbor distance = {estNNdist:#.3g} A")
-            modelScale = ScoreBasedDenoising.originalScale[self.structure] / estNNdist
+            if self.scale is not None:
+                modelScale = (
+                    ScoreBasedDenoising.originalScale[self.structure] / self.scale
+                )
+            else:
+                estNNdist = self.estimateNearestNeighborsDistance(data)
+                print(f"Estimated nearest neighbor distance = {estNNdist:#.3g} A")
+                modelScale = (
+                    ScoreBasedDenoising.originalScale[self.structure] / estNNdist
+                )
 
         denoised_atoms, convergence = yield from self.denoise_snapshot(
             noisy_atoms, model, modelScale