From 5a9f5532d7dbb370b0509adcbc1ba5b2c68c5d5e Mon Sep 17 00:00:00 2001
From: Marcello Sega <m.sega@ucl.ac.uk>
Date: Tue, 6 Aug 2024 12:20:48 +0100
Subject: [PATCH] zero radis handling

---
 pytim/gitim.py            | 3 +++
 pytim/interface.py        | 3 +++
 pytim/itim.py             | 3 +++
 pytim/sanity_check.py     | 5 +++++
 pytim/simple_interface.py | 3 +++
 pytim/willard_chandler.py | 3 +++
 6 files changed, 20 insertions(+)

diff --git a/pytim/gitim.py b/pytim/gitim.py
index 4cf32013..d8879b48 100644
--- a/pytim/gitim.py
+++ b/pytim/gitim.py
@@ -124,6 +124,7 @@ def __init__(self,
                  max_layers=1,
                  radii_dict=None,
                  cluster_cut=None,
+                 include_zero_radius=False,
                  cluster_threshold_density=None,
                  extra_cluster_groups=None,
                  biggest_cluster_only=False,
@@ -155,10 +156,12 @@ def __init__(self,
         self.normal = None
         self.PDB = {}
         self.molecular = molecular
+        self.include_zero_radius = include_zero_radius
         sanity.assign_cluster_params(cluster_cut,
                                      cluster_threshold_density, extra_cluster_groups)
         sanity.check_multiple_layers_options()
         sanity.assign_radii(radii_dict=radii_dict)
+        sanity.prune_zero_radius_atoms()
 
         self._assign_symmetry(symmetry)
         try:
diff --git a/pytim/interface.py b/pytim/interface.py
index e3a9bd0f..fa169c41 100644
--- a/pytim/interface.py
+++ b/pytim/interface.py
@@ -73,6 +73,9 @@ class Interface(object):
     autoassign, _autoassign =\
         _create_property('autoassign',
                          "(bool) assign layers every time a frame changes")
+    include_zero_radius, _include_zero_radius=\
+        _create_property('include_zero_radius',
+                         "(bool) include atoms with zero radius in the analysis (excluded by default)")
     cluster_threshold_density, _cluster_threshold_density =\
         _create_property('cluster_threshold_density',
                          "(float) threshold for the density-based filtering")
diff --git a/pytim/itim.py b/pytim/itim.py
index 1e154d9a..c3b00926 100644
--- a/pytim/itim.py
+++ b/pytim/itim.py
@@ -186,6 +186,7 @@ def __init__(self,
                  max_layers=1,
                  radii_dict=None,
                  cluster_cut=None,
+                 include_zero_radius=False,
                  cluster_threshold_density=None,
                  extra_cluster_groups=None,
                  info=False,
@@ -214,11 +215,13 @@ def __init__(self,
         self.normal = None
         self.PDB = {}
         self.molecular = molecular
+        self.include_zero_radius = include_zero_radius
 
         sanity.assign_cluster_params(cluster_cut,
                                      cluster_threshold_density, extra_cluster_groups)
         sanity.assign_normal(normal)
         sanity.assign_radii(radii_dict=radii_dict)
+        sanity.prune_zero_radius_atoms()
 
         self.grid = None
         self.use_threads = False
diff --git a/pytim/sanity_check.py b/pytim/sanity_check.py
index bda1c1a2..daf79687 100644
--- a/pytim/sanity_check.py
+++ b/pytim/sanity_check.py
@@ -179,6 +179,11 @@ def assign_radii(self, radii_dict=None):
         except BaseException:
             pass
 
+    def prune_zero_radius_atoms(self):
+        if self.interface.include_zero_radius == False:
+            cond = self.interface.analysis_group.radii > 0.0
+            self.interface.analysis_group = self.interface.analysis_group[cond]
+
     def patch_radii_dict(self):
         """ Fix here by hand common problems with radii assignment """
         self.interface.radii_dict['D'] = 0.0
diff --git a/pytim/simple_interface.py b/pytim/simple_interface.py
index cedb7e1a..519d5e6e 100644
--- a/pytim/simple_interface.py
+++ b/pytim/simple_interface.py
@@ -82,6 +82,7 @@ def __init__(self,
                  alpha=1.5,
                  symmetry='generic',
                  normal='z',
+                 include_zero_radius=False,
                  upper=None,
                  lower=None):
 
@@ -89,6 +90,7 @@ def __init__(self,
         self.universe = universe
         self.group = group
         self.alpha = alpha
+        self.include_zero_radius = include_zero_radius
         self.upper = upper
         self.lower = lower
         emptyg = universe.atoms[0:0]
@@ -99,6 +101,7 @@ def __init__(self,
         sanity.assign_universe(universe, group)
         sanity.assign_alpha(alpha)
         sanity.assign_radii(radii_dict=None)
+        sanity.prune_zero_radius_atoms()
         if normal in [0, 1, 2]:
             self.normal = normal
         else:
diff --git a/pytim/willard_chandler.py b/pytim/willard_chandler.py
index 71bff17b..1831a82e 100644
--- a/pytim/willard_chandler.py
+++ b/pytim/willard_chandler.py
@@ -116,6 +116,7 @@ def __init__(self,
                  mesh=2.0,
                  symmetry='spherical',
                  cluster_cut=None,
+                 include_zero_radius=False,
                  cluster_threshold_density=None,
                  extra_cluster_groups=None,
                  centered=False,
@@ -124,6 +125,7 @@ def __init__(self,
                  **kargs):
 
         self.autoassign, self.do_center = autoassign, centered
+        self.include_zero_radius = include_zero_radius
         sanity = SanityCheck(self, warnings=warnings)
         sanity.assign_universe(universe, group)
         sanity.assign_alpha(alpha)
@@ -133,6 +135,7 @@ def __init__(self,
         self.mesh, self.spacing, self.ngrid, self.PDB = mesh, None, None, {}
 
         sanity.assign_radii(radii_dict=radii_dict)
+        sanity.prune_zero_radius_atoms()
 
         sanity.assign_cluster_params(cluster_cut,
                                      cluster_threshold_density, extra_cluster_groups)