From 7666957fa1ed3061d9adb8dba5e98a1a4f77ac4f Mon Sep 17 00:00:00 2001
From: miili <marius.isken@gmail.com>
Date: Sun, 5 Nov 2023 21:01:20 +0100
Subject: [PATCH] fixing tests

---
 lassie/images/__init__.py                     |  7 +++++-
 lassie/search.py                              |  8 -------
 lassie/tracers/fast_marching/fast_marching.py | 23 +++++++++++++------
 .../tracers/fast_marching/velocity_models.py  | 20 ++++++++++------
 4 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/lassie/images/__init__.py b/lassie/images/__init__.py
index ae3080e1..e53dd7cd 100644
--- a/lassie/images/__init__.py
+++ b/lassie/images/__init__.py
@@ -139,11 +139,16 @@ def apply_exponent(self, exponent: float) -> None:
             image.apply_exponent(exponent)
 
     def set_stations(self, stations: Stations) -> None:
-        """Set the images stations."""
+        """Set the images stations.
+
+        Args:
+            stations (Stations): Stations to set.
+        """
         for image in self:
             image.set_stations(stations)
 
     def cumulative_weight(self) -> float:
+        """Get the cumulative weight of all images."""
         return sum(image.weight for image in self)
 
     def snuffle(self) -> None:
diff --git a/lassie/search.py b/lassie/search.py
index 5b159bf8..1234db01 100644
--- a/lassie/search.py
+++ b/lassie/search.py
@@ -638,13 +638,5 @@ async def search(
                 )
 
             detections.append(detection)
-            # detection.plot()
-
-        # plot_octree_movie(octree, semblance, file=Path("/tmp/test.mp4"))
-        if parent.plot_octree_surface:
-            octree.map_semblance(semblance.maximum_node_semblance())
-            parent._plot_octree_surface(
-                octree, time=self.start_time, detections=detections
-            )
 
         return detections, semblance.get_trace()
diff --git a/lassie/tracers/fast_marching/fast_marching.py b/lassie/tracers/fast_marching/fast_marching.py
index 91434f35..d621ab62 100644
--- a/lassie/tracers/fast_marching/fast_marching.py
+++ b/lassie/tracers/fast_marching/fast_marching.py
@@ -87,21 +87,30 @@ def free_cache(self):
 
     def model_post_init(self, __context: Any) -> None:
         grid_spacing = self.grid_spacing
-
-        self._east_coords = np.arange(
+        size_east = self.east_bounds[1] - self.east_bounds[0]
+        size_north = self.north_bounds[1] - self.north_bounds[0]
+        size_depth = self.depth_bounds[1] - self.depth_bounds[0]
+        neast = int(round(size_east / grid_spacing)) + 1
+        nnorth = int(round(size_north / grid_spacing)) + 1
+        ndepth = int(round(size_depth / grid_spacing)) + 1
+
+        self._east_coords = np.linspace(
             self.east_bounds[0],
             self.east_bounds[1],
-            grid_spacing,
+            neast,
+            endpoint=True,
         )
-        self._north_coords = np.arange(
+        self._north_coords = np.linspace(
             self.north_bounds[0],
             self.north_bounds[1],
-            grid_spacing,
+            nnorth,
+            endpoint=True,
         )
-        self._depth_coords = np.arange(
+        self._depth_coords = np.linspace(
             self.depth_bounds[0],
             self.depth_bounds[1],
-            grid_spacing,
+            ndepth,
+            endpoint=True,
         )
 
     @classmethod
diff --git a/lassie/tracers/fast_marching/velocity_models.py b/lassie/tracers/fast_marching/velocity_models.py
index b7996a9a..cede2908 100644
--- a/lassie/tracers/fast_marching/velocity_models.py
+++ b/lassie/tracers/fast_marching/velocity_models.py
@@ -51,21 +51,27 @@ class VelocityModel3D(BaseModel):
 
     def model_post_init(self, __context: Any) -> None:
         grid_spacing = self.grid_spacing
+        neast = int(round(self.east_size / grid_spacing)) + 1
+        nnorth = int(round(self.north_size / grid_spacing)) + 1
+        ndepth = int(round(self.depth_size / grid_spacing)) + 1
 
-        self._east_coords = np.arange(
+        self._east_coords = np.linspace(
             self.east_bounds[0],
             self.east_bounds[1],
-            grid_spacing,
+            neast,
+            endpoint=True,
         )
-        self._north_coords = np.arange(
+        self._north_coords = np.linspace(
             self.north_bounds[0],
             self.north_bounds[1],
-            grid_spacing,
+            nnorth,
+            endpoint=True,
         )
-        self._depth_coords = np.arange(
+        self._depth_coords = np.linspace(
             self.depth_bounds[0],
             self.depth_bounds[1],
-            grid_spacing,
+            ndepth,
+            endpoint=True,
         )
 
         self._velocity_model = np.zeros(
@@ -399,7 +405,7 @@ def north_bounds(self) -> tuple[float, float]:
     @property
     def depth_bounds(self) -> tuple[float, float]:
         """Relative to center location."""
-        return 0.0, self.delta_z * self.nz
+        return 0.0, self.delta_z * (self.nz - 1)
 
     @property
     def center(self) -> Location: