From db6ee27699eae46cf4046207ca6d9260a10ef9ca Mon Sep 17 00:00:00 2001
From: Sylvain Tricot <sylvain.tricot@univ-rennes.fr>
Date: Wed, 26 Feb 2025 11:15:03 +0100
Subject: [PATCH] Add a shape_cluster function

This function is meant to replace the hemispherical_cluster
function dur to some bugs.
---
 src/msspec/utils.py | 102 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 97 insertions(+), 5 deletions(-)

diff --git a/src/msspec/utils.py b/src/msspec/utils.py
index b5f4843..1c37f92 100644
--- a/src/msspec/utils.py
+++ b/src/msspec/utils.py
@@ -19,8 +19,8 @@
 # along with this msspec.  If not, see <http://www.gnu.org/licenses/>.
 #
 # Source file  : src/msspec/utils.py
-# Last modified: Thu, 06 Oct 2022 18:27:24 +0200
-# Committed by : Sylvain Tricot <sylvain.tricot@univ-rennes1.fr> 1665073644 +0200
+# Last modified: Wed, 26 Feb 2025 11:15:03 +0100
+# Committed by : Sylvain Tricot <sylvain.tricot@univ-rennes.fr>
 
 
 """
@@ -468,8 +468,12 @@ def hemispherical_cluster(cluster, emitter_tag=0, emitter_plane=0, diameter=0,
     # the symbol of your emitter
     symbol = cluster[np.where(cluster.get_tags() == emitter_tag)[0][0]].symbol
 
-    assert (diameter != 0 or planes != 0), \
-        "At least one of diameter or planes parameter must be use."
+    if shape.lower() in ('spherical'):
+        assert (diameter != 0 or planes != 0), \
+            "At least one of diameter or planes parameter must be use."
+    elif shape.lower() in ('cylindrical'):
+        assert (diameter != 0 and planes != 0), \
+            "Diameter and planes parameters must be defined for cylindrical shape."
 
     if diameter == 0:
         # calculate the minimal diameter according to the number of planes
@@ -479,6 +483,7 @@ def hemispherical_cluster(cluster, emitter_tag=0, emitter_plane=0, diameter=0,
 
     # number of repetition in each direction
     rep = int(3*min_diameter/min(a, c))
+    #print("rep = ", rep)
 
     # repeat the cluster
     cluster = cluster.repeat((rep, rep, rep))
@@ -542,7 +547,7 @@ def hemispherical_cluster(cluster, emitter_tag=0, emitter_plane=0, diameter=0,
         xplan, yplan = get_xypos(cluster, zplan)
         radius = np.sqrt(xplan**2 + yplan**2 + zplan**2)
 
-        if diameter != 0:
+        if diameter != 0 and shape in ('spherical'):
             assert (radius <= diameter/2), ("The number of planes is too high "
                                             "compared to the diameter.")
         radius = max(radius, diameter/2)
@@ -575,3 +580,90 @@ def hemispherical_cluster(cluster, emitter_tag=0, emitter_plane=0, diameter=0,
     Atoms.translate(cluster, [0, 0, -ze])   # put the emitter in (0,0,0)
 
     return cluster
+
+def shape_cluster(primitive, emitter_tag=0, emitter_plane=0, diameter=0,
+                          planes=0, shape='spherical'):
+
+    """Creates and returns a cluster based on an Atoms object and some
+    parameters.
+
+    :param cluster: the Atoms object used to create the cluster
+    :type  cluster: Atoms object
+    :param emitter_tag: the tag of your emitter
+    :type  emitter_tag: integer
+    :param diameter: the diameter of your cluster in Angströms
+    :type  diameter: float
+    :param planes: the number of planes of your cluster
+    :type  planes: integer
+    :param emitter_plane: the plane where your emitter will be starting by 0
+                          for the first plane
+    :type  emitter_plane: integer
+
+    See :ref:`hemispherical_cluster_faq` for more informations.
+    """
+    # We need the radius of the cluster and the number of planes
+    if shape.lower() in ('ispherical', 'cylindrical'):
+        assert (nplanes != 0 and diameter != 0), "nplanes and diameter cannot be zero for '{}' shape".format(shape)
+    elif shape.lower() in ('spherical'):
+        if diameter <= 0:
+            # find the diameter based on the number of planes
+            assert planes != 0, "planes should be > 0"
+
+
+    n = 3
+    natoms = 0
+    while True:
+        n += 2
+        cluster = primitive.copy()
+        # Repeat the primitive cell
+        cluster = cluster.repeat((n, n, n))
+        center_cluster(cluster)
+
+        # Find the emitter closest to the origin
+        all_tags = cluster.get_tags()
+        are_emitters = all_tags == emitter_tag
+        _ie = np.linalg.norm(cluster[are_emitters].positions, axis=1).argmin()
+        ie = np.nonzero(are_emitters)[0][_ie]
+        # Translate the cluster to this emitter position
+        cluster.translate(-cluster[ie].position)
+        # cut plane at surface and at bottom
+        all_z = np.unique(cluster.positions[:,2])
+        try:
+            zsurf = all_z[all_z >= 0][emitter_plane]
+        except IndexError:
+            # There are not enough planes above the emitter
+            zsurf = all_z.max()
+        try:
+            zbottom = all_z[all_z <= 0][::-1][planes - (emitter_plane+1)]
+        except IndexError:
+            # There are not enough planes below the emitter
+            zbottom = all_z.min()
+        cluster = cut_plane(cluster, z=[zbottom,zsurf])
+        # spherical shape
+        if shape.lower() in ('spherical'):
+            cluster = cut_sphere(cluster, radius=diameter/2, center=(0,0,zsurf))
+        if shape.lower() in ('ispherical'):
+            cluster = cut_sphere(cluster, radius=diameter/2, center=(0,0,0))
+        elif shape.lower() in ('cylindrical'):
+            cluster = cut_cylinder(cluster, radius=diameter/2)
+        else:
+            raise NameError("Unknown shape")
+        cluster.set_cell(primitive.cell)
+        if len(cluster) <= natoms:
+            break
+        else:
+            natoms = len(cluster)
+
+
+    return cluster
+
+
+
+
+
+
+
+
+
+
+