From 9de362803a62464c5a7325ffb1caf25d4cc5c2d4 Mon Sep 17 00:00:00 2001
From: Andrew Nolan <anolan@lanl.gov>
Date: Fri, 20 Dec 2024 14:31:27 -0700
Subject: [PATCH] Add support for wrapping planar periodic meshes.

Further testing is need to move the spherical wrapping into the
new method used by planar periodic meshes.
---
 mosaic/descriptor.py | 56 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 53 insertions(+), 3 deletions(-)

diff --git a/mosaic/descriptor.py b/mosaic/descriptor.py
index c58705c..1a2123f 100644
--- a/mosaic/descriptor.py
+++ b/mosaic/descriptor.py
@@ -295,7 +295,7 @@ def cell_patches(self) -> ndarray:
         patch. Nodes are ordered counter clockwise around the cell center.
         """
         patches = _compute_cell_patches(self.ds)
-        patches = self._fix_antimeridian(patches, "Cell")
+        patches = self._wrap_patches(patches, "Cell")
         return patches
 
     @cached_property
@@ -314,7 +314,7 @@ def edge_patches(self) -> ndarray:
         corresponding node will be collapsed to the edge coordinate.
         """
         patches = _compute_edge_patches(self.ds)
-        patches = self._fix_antimeridian(patches, "Edge")
+        patches = self._wrap_patches(patches, "Edge")
         return patches
 
     @cached_property
@@ -340,7 +340,7 @@ def vertex_patches(self) -> ndarray:
         position.
         """
         patches = _compute_vertex_patches(self.ds)
-        patches = self._fix_antimeridian(patches, "Vertex")
+        patches = self._wrap_patches(patches, "Vertex")
         return patches
 
     def _transform_coordinates(self, projection, transform):
@@ -355,6 +355,56 @@ def _transform_coordinates(self, projection, transform):
             self.ds[f"x{loc}"].values = transformed_coords[:, 0]
             self.ds[f"y{loc}"].values = transformed_coords[:, 1]
 
+    def _wrap_patches(self, patches, loc):
+        """Wrap patches for spherical and planar-periodic meshes
+
+        """
+
+        def _find_boundary_patches(patches, loc, coord):
+            """
+            Find the patches that cross the periodic boundary and what
+            direction they cross the boundary (i.e. their ``sign``). This
+            method assumes the patch centroids are not periodic
+            """
+            # get axis index we are inquiring over
+            axis = 0 if coord == "x" else 1
+            # get requested coordinate of patch centroids
+            center = self.ds[f"{coord}{loc.title()}"].values.reshape(-1, 1)
+            # get difference b/w centroid and nodes of patches
+            diff = patches[..., axis] - center
+
+            #
+            if self.__getattribute__(f"{coord}_period"):
+                period = self.__getattribute__(f"{coord}_period")
+
+            mask = np.abs(diff) > np.abs(period) / (2. * np.sqrt(2.))
+            sign = np.sign(diff)
+
+            return mask, sign
+
+        def _wrap_1D(patches, mask, sign, axis, period):
+            """Correct patch periodicity along a single dimension"""
+            patches[..., axis][mask] -= np.sign(sign[mask]) * period
+            return patches
+
+        if self.x_period:
+            # find the patch that are periodic in x-direction
+            x_mask, x_sign = _find_boundary_patches(patches, loc, "x")
+            # using the sign of the difference correct patches x coordinate
+            patches = _wrap_1D(patches, x_mask, x_sign, 0, self.x_period)
+
+        if self.y_period:
+            # find the patch that are periodic in y-direction
+            y_mask, y_sign = _find_boundary_patches(patches, loc, "y")
+            # using the sign of the difference correct patches y coordinate
+            patches = _wrap_1D(patches, y_mask, y_sign, 1, self.y_period)
+
+        if self.is_spherical:
+            # call current spherical wrapping function for now
+            patches = self._fix_antimeridian(patches, loc)
+
+        return patches
+
     def _fix_antimeridian(self, patches, loc, projection=None):
         """Correct vertices of patches that cross the antimeridian.