Skip to content

Commit a58e731

Browse files
wpbonelliwpbonelli
committed
refactor(particle tracking output): consolidation/cleanup, support prt files
1 parent 9e87acd commit a58e731

File tree

5 files changed

+734
-417
lines changed

5 files changed

+734
-417
lines changed

‎autotest/test_plotutil.py

+117-65
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,17 @@
33
import pytest
44

55
from flopy.plot.plotutil import (
6-
MP7_ENDPOINT_DTYPE,
7-
MP7_PATHLINE_DTYPE,
8-
PRT_PATHLINE_DTYPE,
96
to_mp7_endpoints,
107
to_mp7_pathlines,
118
to_prt_pathlines,
129
)
10+
from flopy.utils.modpathfile import (
11+
EndpointFile as MpEndpointFile,
12+
)
13+
from flopy.utils.modpathfile import (
14+
PathlineFile as MpPathlineFile,
15+
)
16+
from flopy.utils.prtfile import PathlineFile as PrtPathlineFile
1317

1418
PRT_TEST_PATHLINES = pd.DataFrame.from_records(
1519
[
@@ -194,7 +198,7 @@
194198
"PRP000000001", # name
195199
],
196200
],
197-
columns=PRT_PATHLINE_DTYPE.names,
201+
columns=PrtPathlineFile.dtypes["base"].names,
198202
)
199203
MP7_TEST_PATHLINES = pd.DataFrame.from_records(
200204
[
@@ -233,7 +237,7 @@
233237
1, # timestep
234238
],
235239
],
236-
columns=MP7_PATHLINE_DTYPE.names,
240+
columns=MpPathlineFile.dtypes[7].names,
237241
)
238242
MP7_TEST_ENDPOINTS = pd.DataFrame.from_records(
239243
[
@@ -322,114 +326,152 @@
322326
2, # cellface
323327
],
324328
],
325-
columns=MP7_ENDPOINT_DTYPE.names,
329+
columns=MpEndpointFile.dtypes[7].names,
326330
)
327331

328332

329333
@pytest.mark.parametrize("dataframe", [True, False])
330-
def test_to_mp7_pathlines(dataframe):
331-
prt_pls = (
332-
PRT_TEST_PATHLINES
333-
if dataframe
334-
else PRT_TEST_PATHLINES.to_records(index=False)
335-
)
336-
mp7_pls = to_mp7_pathlines(prt_pls)
334+
@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"])
335+
def test_to_mp7_pathlines(dataframe, source):
336+
if source == "prt":
337+
pls = (
338+
PRT_TEST_PATHLINES
339+
if dataframe
340+
else PRT_TEST_PATHLINES.to_records(index=False)
341+
)
342+
elif source == "mp3":
343+
pass
344+
elif source == "mp5":
345+
pass
346+
elif source == "mp7":
347+
pass
348+
mp7_pls = to_mp7_pathlines(pls)
337349
assert (
338-
type(prt_pls)
350+
type(pls)
339351
== type(mp7_pls)
340352
== (pd.DataFrame if dataframe else np.recarray)
341353
)
342354
assert len(mp7_pls) == 10
343355
assert set(
344356
dict(mp7_pls.dtypes).keys() if dataframe else mp7_pls.dtype.names
345-
) == set(MP7_PATHLINE_DTYPE.names)
357+
) == set(MpPathlineFile.dtypes[7].names)
346358

347359

348360
@pytest.mark.parametrize("dataframe", [True, False])
349-
def test_to_mp7_pathlines_empty(dataframe):
350-
mp7_pls = to_mp7_pathlines(
351-
pd.DataFrame.from_records([], columns=PRT_PATHLINE_DTYPE.names)
352-
if dataframe
353-
else np.recarray((0,), dtype=PRT_PATHLINE_DTYPE)
354-
)
355-
assert mp7_pls.empty if dataframe else mp7_pls.size == 0
361+
@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"])
362+
def test_to_mp7_pathlines_empty(dataframe, source):
363+
if source == "prt":
364+
pls = to_mp7_pathlines(
365+
pd.DataFrame.from_records(
366+
[], columns=PrtPathlineFile.dtypes["base"].names
367+
)
368+
if dataframe
369+
else np.recarray((0,), dtype=PrtPathlineFile.dtypes["base"])
370+
)
371+
elif source == "mp3":
372+
pass
373+
elif source == "mp5":
374+
pass
375+
elif source == "mp7":
376+
pass
377+
assert pls.empty if dataframe else pls.size == 0
356378
if dataframe:
357-
mp7_pls = mp7_pls.to_records(index=False)
358-
assert mp7_pls.dtype == MP7_PATHLINE_DTYPE
379+
pls = pls.to_records(index=False)
380+
assert pls.dtype == MpPathlineFile.dtypes[7]
359381

360382

361383
@pytest.mark.parametrize("dataframe", [True, False])
362384
def test_to_mp7_pathlines_noop(dataframe):
363-
prt_pls = (
385+
pls = (
364386
MP7_TEST_PATHLINES
365387
if dataframe
366388
else MP7_TEST_PATHLINES.to_records(index=False)
367389
)
368-
mp7_pls = to_mp7_pathlines(prt_pls)
390+
mp7_pls = to_mp7_pathlines(pls)
369391
assert (
370-
type(prt_pls)
392+
type(pls)
371393
== type(mp7_pls)
372394
== (pd.DataFrame if dataframe else np.recarray)
373395
)
374396
assert len(mp7_pls) == 2
375397
assert set(
376398
dict(mp7_pls.dtypes).keys() if dataframe else mp7_pls.dtype.names
377-
) == set(MP7_PATHLINE_DTYPE.names)
399+
) == set(MpPathlineFile.dtypes[7].names)
378400
assert np.array_equal(
379401
mp7_pls if dataframe else pd.DataFrame(mp7_pls), MP7_TEST_PATHLINES
380402
)
381403

382404

383405
@pytest.mark.parametrize("dataframe", [True, False])
384-
def test_to_mp7_endpoints(dataframe):
385-
mp7_eps = to_mp7_endpoints(
386-
PRT_TEST_PATHLINES
387-
if dataframe
388-
else PRT_TEST_PATHLINES.to_records(index=False)
389-
)
390-
assert len(mp7_eps) == 1
391-
assert np.isclose(mp7_eps.time[0], PRT_TEST_PATHLINES.t.max())
406+
@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"])
407+
def test_to_mp7_endpoints(dataframe, source):
408+
if source == "prt":
409+
eps = to_mp7_endpoints(
410+
PRT_TEST_PATHLINES
411+
if dataframe
412+
else PRT_TEST_PATHLINES.to_records(index=False)
413+
)
414+
elif source == "mp3":
415+
pass
416+
elif source == "mp5":
417+
pass
418+
elif source == "mp6":
419+
pass
420+
assert len(eps) == 1
421+
assert np.isclose(eps.time[0], PRT_TEST_PATHLINES.t.max())
392422
assert set(
393-
dict(mp7_eps.dtypes).keys() if dataframe else mp7_eps.dtype.names
394-
) == set(MP7_ENDPOINT_DTYPE.names)
423+
dict(eps.dtypes).keys() if dataframe else eps.dtype.names
424+
) == set(MpEndpointFile.dtypes[7].names)
395425

396426

397427
@pytest.mark.parametrize("dataframe", [True, False])
398-
def test_to_mp7_endpoints_empty(dataframe):
399-
mp7_eps = to_mp7_endpoints(
400-
pd.DataFrame.from_records([], columns=PRT_PATHLINE_DTYPE.names)
428+
@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"])
429+
def test_to_mp7_endpoints_empty(dataframe, source):
430+
eps = to_mp7_endpoints(
431+
pd.DataFrame.from_records(
432+
[], columns=PrtPathlineFile.dtypes["base"].names
433+
)
401434
if dataframe
402-
else np.recarray((0,), dtype=PRT_PATHLINE_DTYPE)
435+
else np.recarray((0,), dtype=PrtPathlineFile.dtypes["base"])
403436
)
404-
assert mp7_eps.empty if dataframe else mp7_eps.size == 0
437+
assert eps.empty if dataframe else eps.size == 0
405438
if dataframe:
406-
mp7_eps = mp7_eps.to_records(index=False)
407-
assert mp7_eps.dtype == MP7_ENDPOINT_DTYPE
439+
eps = eps.to_records(index=False)
440+
assert eps.dtype == MpEndpointFile.dtypes[7]
408441

409442

410443
@pytest.mark.parametrize("dataframe", [True, False])
411444
def test_to_mp7_endpoints_noop(dataframe):
412445
"""Test a recarray or dataframe which already contains MP7 endpoint data"""
413-
mp7_eps = to_mp7_endpoints(
446+
eps = to_mp7_endpoints(
414447
MP7_TEST_ENDPOINTS
415448
if dataframe
416449
else MP7_TEST_ENDPOINTS.to_records(index=False)
417450
)
418451
assert np.array_equal(
419-
mp7_eps if dataframe else pd.DataFrame(mp7_eps), MP7_TEST_ENDPOINTS
452+
eps if dataframe else pd.DataFrame(eps), MP7_TEST_ENDPOINTS
420453
)
421454

422455

423456
@pytest.mark.parametrize("dataframe", [True, False])
424-
def test_to_prt_pathlines_roundtrip(dataframe):
425-
mp7_pls = to_mp7_pathlines(
426-
PRT_TEST_PATHLINES
427-
if dataframe
428-
else PRT_TEST_PATHLINES.to_records(index=False)
429-
)
430-
prt_pls = to_prt_pathlines(mp7_pls)
457+
@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"])
458+
def test_to_prt_pathlines_roundtrip(dataframe, source):
459+
if source == "prt":
460+
pls = to_mp7_pathlines(
461+
PRT_TEST_PATHLINES
462+
if dataframe
463+
else PRT_TEST_PATHLINES.to_records(index=False)
464+
)
465+
elif source == "mp3":
466+
pass
467+
elif source == "mp5":
468+
pass
469+
elif source == "mp6":
470+
pass
471+
prt_pls = to_prt_pathlines(pls)
431472
if not dataframe:
432473
prt_pls = pd.DataFrame(prt_pls)
474+
# import pdb; pdb.set_trace()
433475
assert np.allclose(
434476
PRT_TEST_PATHLINES.drop(
435477
["imdl", "iprp", "irpt", "name", "istatus", "ireason"],
@@ -443,15 +485,25 @@ def test_to_prt_pathlines_roundtrip(dataframe):
443485

444486

445487
@pytest.mark.parametrize("dataframe", [True, False])
446-
def test_to_prt_pathlines_roundtrip_empty(dataframe):
447-
mp7_pls = to_mp7_pathlines(
448-
pd.DataFrame.from_records([], columns=PRT_PATHLINE_DTYPE.names)
449-
if dataframe
450-
else np.recarray((0,), dtype=PRT_PATHLINE_DTYPE)
451-
)
452-
prt_pls = to_prt_pathlines(mp7_pls)
453-
assert mp7_pls.empty if dataframe else mp7_pls.size == 0
454-
assert prt_pls.empty if dataframe else mp7_pls.size == 0
488+
@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"])
489+
def test_to_prt_pathlines_roundtrip_empty(dataframe, source):
490+
if source == "prt":
491+
pls = to_mp7_pathlines(
492+
pd.DataFrame.from_records(
493+
[], columns=PrtPathlineFile.dtypes["base"].names
494+
)
495+
if dataframe
496+
else np.recarray((0,), dtype=PrtPathlineFile.dtypes["base"])
497+
)
498+
elif source == "mp3":
499+
pass
500+
elif source == "mp5":
501+
pass
502+
elif source == "mp6":
503+
pass
504+
prt_pls = to_prt_pathlines(pls)
505+
assert pls.empty if dataframe else pls.size == 0
506+
assert prt_pls.empty if dataframe else pls.size == 0
455507
assert set(
456-
dict(mp7_pls.dtypes).keys() if dataframe else mp7_pls.dtype.names
457-
) == set(MP7_PATHLINE_DTYPE.names)
508+
dict(pls.dtypes).keys() if dataframe else pls.dtype.names
509+
) == set(MpPathlineFile.dtypes[7].names)

0 commit comments

Comments
 (0)