fix issue, expand tests

modflowpy · Jan 25, 2024 · 686074d · 686074d
1 parent 7dfc968
commit 686074d
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 2 deletions.
diff --git a/autotest/test_particledata.py b/autotest/test_particledata.py
@@ -129,6 +129,46 @@ def test_particledata_unstructured_ctor_with_partlocs_as_ndarray():
     )
 
 
+def test_particledata_unstructured_ctor_with_partlocs_as_list():
+    locs = [0, 1, 2]
+    data = ParticleData(partlocs=locs, structured=False)
+
+    assert data.particlecount == 3
+    assert data.dtype == unstructured_dtype
+    assert isinstance(data.particledata, pd.DataFrame)
+    assert np.array_equal(
+        data.particledata.to_records(index=False),
+        np.core.records.fromrecords(
+            [
+                (0, 0.5, 0.5, 0.5, 0.0, 0),
+                (1, 0.5, 0.5, 0.5, 0.0, 0),
+                (2, 0.5, 0.5, 0.5, 0.0, 0),
+            ],
+            dtype=unstructured_dtype,
+        ),
+    )
+
+
+def test_particledata_unstructured_ctor_with_partlocs_as_ndarray():
+    locs = np.array([0, 1, 2])
+    data = ParticleData(partlocs=locs, structured=False)
+
+    assert data.particlecount == 3
+    assert data.dtype == unstructured_dtype
+    assert isinstance(data.particledata, pd.DataFrame)
+    assert np.array_equal(
+        data.particledata.to_records(index=False),
+        np.core.records.fromrecords(
+            [
+                (0, 0.5, 0.5, 0.5, 0.0, 0),
+                (1, 0.5, 0.5, 0.5, 0.0, 0),
+                (2, 0.5, 0.5, 0.5, 0.0, 0),
+            ],
+            dtype=unstructured_dtype,
+        ),
+    )
+
+
 def test_particledata_structured_ctor_with_partlocs_as_list_of_lists():
     locs = [list(p) for p in [(0, 1, 1), (0, 1, 2)]]
     data = ParticleData(partlocs=locs, structured=True)

diff --git a/flopy/modpath/mp7particledata.py b/flopy/modpath/mp7particledata.py
@@ -166,12 +166,17 @@ def __init__(
                         "one entry".format(self.name)
                     )
 
-            # convert partlocs composed of a lists/tuples of lists/tuples
-            # to a numpy array
+            # convert partlocs to numpy array with proper dtype
+            partlocs = np.array(partlocs)
+            if len(partlocs.shape) == 1:
+                partlocs = partlocs.reshape(len(partlocs), 1)
             partlocs = unstructured_to_structured(
                 np.array(partlocs), dtype=dtype
             )
         elif isinstance(partlocs, np.ndarray):
+            # reshape and convert dtype if needed
+            if len(partlocs.shape) == 1:
+                partlocs = partlocs.reshape(len(partlocs), 1)
             dtypein = partlocs.dtype
             if dtypein != dtype:
                 partlocs = unstructured_to_structured(partlocs, dtype=dtype)