Remove unused variables

eulerpi/examples/heat/heat.py

@@ -104,17 +104,30 @@ def forward(cls, param: np.ndarray) -> np.ndarray:
 
         solution = cls.perform_simulation(kappa=param)
 
-        # the indices of the wanted evaluation points
-        evaluation_indices = jnp.multiply(
-            cls.evaluation_points,
-            jnp.array(solution.shape),
-        ).astype(int)
-
-        solution_at_evaluation_points = jnp.array(
-            solution[
-                evaluation_indices[:, 0],
-                evaluation_indices[:, 1],
-            ]
+        # linearly interpolate the solution at the evaluation points, use own interpolation function as jax doesn't support scipy.interpolate.interp2d and doesn't provide a 2d interpolation function
+        x = jnp.linspace(0, cls.plate_length[0], cls.num_grid_points)
+        y = jnp.linspace(0, cls.plate_length[1], cls.num_grid_points)
+
+        # compute the indices between which the evaluation points lie
+        x_indices = jnp.searchsorted(x, cls.evaluation_points[:, 0]) - 1
+        y_indices = jnp.searchsorted(y, cls.evaluation_points[:, 1]) - 1
+        dx = cls.plate_length[0] / cls.num_grid_points
+        dy = cls.plate_length[1] / cls.num_grid_points
+
+        # interpolate the solution at the evaluation points
+        solution_at_evaluation_points = (1 / (dx * dy)) * (
+            solution[x_indices, y_indices]
+            * (x[x_indices + 1] - cls.evaluation_points[:, 0])
+            * (y[y_indices + 1] - cls.evaluation_points[:, 1])
+            + solution[x_indices + 1, y_indices]
+            * (cls.evaluation_points[:, 0] - x[x_indices])
+            * (y[y_indices + 1] - cls.evaluation_points[:, 1])
+            + solution[x_indices, y_indices + 1]
+            * (x[x_indices + 1] - cls.evaluation_points[:, 0])
+            * (cls.evaluation_points[:, 1] - y[y_indices])
+            + solution[x_indices + 1, y_indices + 1]
+            * (cls.evaluation_points[:, 0] - x[x_indices])
+            * (cls.evaluation_points[:, 1] - y[y_indices])
         )
         return solution_at_evaluation_points
 
@@ -196,6 +209,7 @@ class HeatArtificial(Heat, ArtificialModelInterface):
 
     CENTRAL_PARAM = np.array([1.5, 1.5, 0.5])
     PARAM_LIMITS = np.array([[1.0, 2.0], [1.0, 2.0], [0.0, 1.0]])
+    num_grid_points = 20
 
     def __init__(
         self,