New geom.at_z method

Author: amorison

stagpy/_step.py

@@ -163,6 +163,14 @@ def threed(self):
         """Whether the grid is 3 dimensional."""
         return not self.twod
 
+    def at_z(self, zval):
+        """Return iz closest to given zval position.
+
+        This is different than radial position in spherical geometry.
+        """
+        return np.argmin(np.abs(self.z_coord - zval))
+
+
     def __getattr__(self, attr):
         # provide nDtot, D_coord, D_mesh and nbtot
         # with D = x, y, z or t, p, r