diff --git a/dataset/point_source/simple/point_dataset_positions_only.json b/dataset/point_source/simple/point_dataset_positions_only.json
index fac29cd..3b2cc93 100644
--- a/dataset/point_source/simple/point_dataset_positions_only.json
+++ b/dataset/point_source/simple/point_dataset_positions_only.json
@@ -3,6 +3,22 @@
     "class_path": "autolens.point.dataset.PointDataset",
     "arguments": {
         "name": "point_0",
+        "fluxes_noise_map": null,
+        "positions_noise_map": {
+            "type": "instance",
+            "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular",
+            "arguments": {
+                "values": {
+                    "type": "ndarray",
+                    "array": [
+                        0.005,
+                        0.005
+                    ],
+                    "dtype": "float64"
+                }
+            }
+        },
+        "time_delays": null,
         "positions": {
             "type": "instance",
             "class_path": "autoarray.structures.grids.irregular_2d.Grid2DIrregular",
@@ -23,23 +39,8 @@
                 }
             }
         },
-        "time_delays_noise_map": null,
         "fluxes": null,
-        "positions_noise_map": {
-            "type": "instance",
-            "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular",
-            "arguments": {
-                "values": {
-                    "type": "ndarray",
-                    "array": [
-                        0.2,
-                        0.2
-                    ],
-                    "dtype": "float64"
-                }
-            }
-        },
-        "time_delays": null,
-        "fluxes_noise_map": null
+        "redshift": null,
+        "time_delays_noise_map": null
     }
 }
\ No newline at end of file
diff --git a/dataset/point_source/simple/tracer.json b/dataset/point_source/simple/tracer.json
index 99344b8..00b6f2e 100644
--- a/dataset/point_source/simple/tracer.json
+++ b/dataset/point_source/simple/tracer.json
@@ -2,11 +2,6 @@
     "type": "instance",
     "class_path": "autolens.lens.tracer.Tracer",
     "arguments": {
-        "cosmology": {
-            "type": "instance",
-            "class_path": "autogalaxy.cosmology.model.Planck15",
-            "arguments": {}
-        },
         "galaxies": {
             "type": "list",
             "values": [
@@ -48,6 +43,11 @@
                     }
                 }
             ]
+        },
+        "cosmology": {
+            "type": "instance",
+            "class_path": "autogalaxy.cosmology.model.Planck15",
+            "arguments": {}
         }
     }
 }
\ No newline at end of file
diff --git a/scripts/jax_likelihood_functions/point_source/image_plane.py b/scripts/jax_likelihood_functions/point_source/image_plane.py
index f084475..bf4a19d 100644
--- a/scripts/jax_likelihood_functions/point_source/image_plane.py
+++ b/scripts/jax_likelihood_functions/point_source/image_plane.py
@@ -132,7 +132,7 @@
 
 np.testing.assert_allclose(
     np.array(result),
-    1.313508,
+    -83.38049778,
     rtol=1e-4,
     err_msg="point_source/image_plane: JAX vmap likelihood mismatch",
 )
diff --git a/scripts/jax_likelihood_functions/point_source/point.py b/scripts/jax_likelihood_functions/point_source/point.py
index 59dbf66..8bc47b7 100644
--- a/scripts/jax_likelihood_functions/point_source/point.py
+++ b/scripts/jax_likelihood_functions/point_source/point.py
@@ -233,7 +233,7 @@
 
 np.testing.assert_allclose(
     np.array(result),
-    1.313508,
+    -83.38049778,
     rtol=1e-4,
     err_msg="point: JAX vmap likelihood mismatch",
 )
diff --git a/scripts/jax_likelihood_functions/point_source/source_plane.py b/scripts/jax_likelihood_functions/point_source/source_plane.py
index b8b1495..48a7c24 100644
--- a/scripts/jax_likelihood_functions/point_source/source_plane.py
+++ b/scripts/jax_likelihood_functions/point_source/source_plane.py
@@ -143,7 +143,7 @@
 
 # Regression constant filled in on first run against the workspace_test
 # seeded PointDataset (``scripts/jax_likelihood_functions/point_source/simulator.py``).
-EXPECTED_VMAP_LOG_LIKELIHOOD_SOURCE_PLANE = -199.1555813
+EXPECTED_VMAP_LOG_LIKELIHOOD_SOURCE_PLANE = -331481.25978149
 
 np.testing.assert_allclose(
     np.array(result),
@@ -180,7 +180,7 @@
 log_likelihood_np = float(fit_np.log_likelihood)
 print("NumPy fit.log_likelihood:", log_likelihood_np)
 
-EXPECTED_EAGER_LOG_LIKELIHOOD_SOURCE_PLANE = -199.1555813
+EXPECTED_EAGER_LOG_LIKELIHOOD_SOURCE_PLANE = -331481.26508536364
 
 np.testing.assert_allclose(
     log_likelihood_np,