diff --git a/autotest/test_mp6.py b/autotest/test_mp6.py index 73935a1d4..11e8b9034 100644 --- a/autotest/test_mp6.py +++ b/autotest/test_mp6.py @@ -308,9 +308,9 @@ def test_loadtxt(function_tmpdir, mp6_test_path): pthfile = function_tmpdir / "EXAMPLE-3.pathline" pthld = PathlineFile(pthfile) - ra = loadtxt(pthfile, delimiter=" ", skiprows=3, dtype=pthld.dtype) + ra = loadtxt(pthfile, delimiter=" ", skiprows=3, dtype=pthld._dtype) ra2 = loadtxt( - pthfile, delimiter=" ", skiprows=3, dtype=pthld.dtype, use_pandas=False + pthfile, delimiter=" ", skiprows=3, dtype=pthld._dtype, use_pandas=False ) assert np.array_equal(ra, ra2) diff --git a/autotest/test_plotutil.py b/autotest/test_plotutil.py index bc5fbb05a..9ed040c54 100644 --- a/autotest/test_plotutil.py +++ b/autotest/test_plotutil.py @@ -1,483 +1,393 @@ +from pathlib import Path + import numpy as np import pandas as pd import pytest +import flopy from flopy.plot.plotutil import ( to_mp7_endpoints, to_mp7_pathlines, to_prt_pathlines, ) -from flopy.utils.modpathfile import ( - EndpointFile as MpEndpointFile, -) -from flopy.utils.modpathfile import ( - PathlineFile as MpPathlineFile, -) +from flopy.utils.modpathfile import EndpointFile as MpEndpointFile +from flopy.utils.modpathfile import PathlineFile as MpPathlineFile from flopy.utils.prtfile import PathlineFile as PrtPathlineFile -PRT_TEST_PATHLINES = pd.DataFrame.from_records( - [ - [ - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 0, - 0, - 0, - 0.0, - 0.000000, - 0.100000, - 9.1, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 1, - 0, - 1, - 1, - 0.0, - 0.063460, - 0.111111, - 9.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 11, - 0, - 1, - 1, - 0.0, - 0.830431, - 0.184020, - 8.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 21, - 0, - 1, - 1, - 0.0, - 2.026390, - 0.267596, - 7.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 31, - 0, - 1, - 1, - 0.0, - 3.704265, - 0.360604, - 6.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 60, - 0, - 1, - 1, - 0.0, - 39.087992, - 9.639587, - 4.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 70, - 0, - 1, - 1, - 0.0, - 40.765791, - 9.732597, - 3.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 80, - 0, - 1, - 1, - 0.0, - 41.961755, - 9.816110, - 2.0, - 0.5, - "PRP000000001", - ], - [ - 1, - 1, - 1, - 1, - 1, - 1, - 90, - 0, - 1, - 1, - 0.0, - 42.728752, - 9.888968, - 1.0, - 0.5, - "PRP000000001", - ], - [ - 1, # kper - 1, # kstp - 1, # imdl - 1, # iprp - 1, # irpt - 1, # ilay - 100, # icell - 0, # izone - 5, # istatus - 3, # ireason - 0.0, # trelease - 42.728752, # t - 9.888968, # x - 1.0, # y - 0.5, # z - "PRP000000001", # name - ], - ], - columns=PrtPathlineFile.dtypes["base"].names, -) -MP7_TEST_PATHLINES = pd.DataFrame.from_records( - [ - [ - 1, # particleid - 1, # particlegroup - 1, # sequencenumber - 1, # particleidloc - 0.0, # time - 1.0, # x - 2.0, # y - 3.0, # z - 1, # k - 1, # node - 0.1, # xloc - 0.1, # yloc - 0.1, # zloc - 1, # stressperiod - 1, # timestep - ], - [ - 1, - 1, - 1, - 1, - 1.0, # time - 2.0, # x - 3.0, # y - 4.0, # z - 2, # k - 2, # node - 0.9, # xloc - 0.9, # yloc - 0.9, # zloc - 1, # stressperiod - 1, # timestep - ], - ], - columns=MpPathlineFile.dtypes[7].names, -) -MP7_TEST_ENDPOINTS = pd.DataFrame.from_records( - [ - [ - 1, # particleid - 1, # particlegroup - 1, # particleidloc - 2, # status (terminated at boundary face) - 0.0, # time0 - 1.0, # time - 1, # node0 - 1, # k0 - 0.1, # xloc0 - 0.1, # yloc0 - 0.1, # zloc0 - 0.0, # x0 - 1.0, # y0 - 2.0, # z0 - 1, # zone0 - 1, # initialcellface - 5, # node - 2, # k - 0.9, # xloc - 0.9, # yloc - 0.9, # zloc - 10.0, # x - 11.0, # y - 12.0, # z - 2, # zone - 2, # cellface + +nlay = 1 +nrow = 10 +ncol = 10 +top = 1.0 +botm = [0.0] +nper = 1 +perlen = 1.0 +nstp = 1 +tsmult = 1.0 +porosity = 0.1 + + +def get_partdata(grid, rpts): + """ + Make a flopy.modpath.ParticleData from the given grid and release points. + """ + + if grid.grid_type == "structured": + return flopy.modpath.ParticleData( + partlocs=[grid.get_lrc(p[0])[0] for p in rpts], + structured=True, + localx=[p[1] for p in rpts], + localy=[p[2] for p in rpts], + localz=[p[3] for p in rpts], + timeoffset=0, + drape=0, + ) + else: + return flopy.modpath.ParticleData( + partlocs=[p[0] for p in rpts], + structured=False, + localx=[p[1] for p in rpts], + localy=[p[2] for p in rpts], + localz=[p[3] for p in rpts], + timeoffset=0, + drape=0, + ) + + +@pytest.fixture +def gwf_sim(function_tmpdir): + gwf_ws = function_tmpdir / "gwf" + gwf_name = "plotutil_gwf" + + # create simulation + sim = flopy.mf6.MFSimulation( + sim_name=gwf_name, + exe_name="mf6", + version="mf6", + sim_ws=gwf_ws, + ) + + # create tdis package + flopy.mf6.modflow.mftdis.ModflowTdis( + sim, + pname="tdis", + time_units="DAYS", + nper=nper, + perioddata=[ + (perlen, nstp, tsmult) ], - [ - 2, # particleid - 1, # particlegroup - 2, # particleidloc - 2, # status (terminated at boundary face) - 0.0, # time0 - 2.0, # time - 1, # node0 - 1, # k0 - 0.1, # xloc0 - 0.1, # yloc0 - 0.1, # zloc0 - 0.0, # x0 - 1.0, # y0 - 2.0, # z0 - 1, # zone0 - 1, # initialcellface - 5, # node - 2, # k - 0.9, # xloc - 0.9, # yloc - 0.9, # zloc - 10.0, # x - 11.0, # y - 12.0, # z - 2, # zone - 2, # cellface + ) + + # create gwf model + gwf = flopy.mf6.ModflowGwf(sim, modelname=gwf_name, save_flows=True) + + # create gwf discretization + flopy.mf6.modflow.mfgwfdis.ModflowGwfdis( + gwf, + pname="dis", + nlay=nlay, + nrow=nrow, + ncol=ncol, + ) + + # create gwf initial conditions package + flopy.mf6.modflow.mfgwfic.ModflowGwfic(gwf, pname="ic") + + # create gwf node property flow package + flopy.mf6.modflow.mfgwfnpf.ModflowGwfnpf( + gwf, + pname="npf", + save_saturation=True, + save_specific_discharge=True, + ) + + # create gwf chd package + spd = { + 0: [[(0, 0, 0), 1.0, 1.0], [(0, 9, 9), 0.0, 0.0]], + 1: [[(0, 0, 0), 0.0, 0.0], [(0, 9, 9), 1.0, 2.0]], + } + chd = flopy.mf6.ModflowGwfchd( + gwf, + pname="CHD-1", + stress_period_data=spd, + auxiliary=["concentration"], + ) + + # create gwf output control package + # output file names + gwf_budget_file = f"{gwf_name}.bud" + gwf_head_file = f"{gwf_name}.hds" + oc = flopy.mf6.ModflowGwfoc( + gwf, + budget_filerecord=gwf_budget_file, + head_filerecord=gwf_head_file, + saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")], + ) + + # create iterative model solution for gwf model + ims = flopy.mf6.ModflowIms(sim) + + return sim + + +@pytest.fixture +def mp7_sim(gwf_sim, function_tmpdir): + gwf = gwf_sim.get_model() + mp7_ws = function_tmpdir / "mp7" + releasepts_mp7 = [ + # node number, localx, localy, localz + (0, float(f"0.{i + 1}"), float(f"0.{i + 1}"), 0.5) + for i in range(9) + ] + + partdata = get_partdata(gwf.modelgrid, releasepts_mp7) + mp7_name = "plotutil_mp7" + pg = flopy.modpath.ParticleGroup( + particlegroupname="G1", + particledata=partdata, + filename=f"{mp7_name}.sloc", + ) + mp = flopy.modpath.Modpath7( + modelname=mp7_name, + flowmodel=gwf, + exe_name="mp7", + model_ws=mp7_ws, + headfilename=f"{gwf.name}.hds", + budgetfilename=f"{gwf.name}.bud", + ) + mpbas = flopy.modpath.Modpath7Bas( + mp, + porosity=porosity, + ) + mpsim = flopy.modpath.Modpath7Sim( + mp, + simulationtype="pathline", + trackingdirection="forward", + budgetoutputoption="summary", + stoptimeoption="extend", + particlegroups=[pg], + ) + + return mp + + +@pytest.fixture +def prt_sim(function_tmpdir): + gwf_ws = function_tmpdir / "gwf" + prt_ws = function_tmpdir / "prt" + prt_name = "plotutil_prt" + gwf_name = "plotutil_gwf" + releasepts_prt = [ + # particle index, k, i, j, x, y, z + [i, 0, 0, 0, float(f"0.{i + 1}"), float(f"9.{i + 1}"), 0.5] + for i in range(9) + ] + + # create simulation + sim = flopy.mf6.MFSimulation( + sim_name=prt_name, + exe_name="mf6", + version="mf6", + sim_ws=prt_ws, + ) + + # create tdis package + flopy.mf6.modflow.mftdis.ModflowTdis( + sim, + pname="tdis", + time_units="DAYS", + nper=nper, + perioddata=[ + ( + perlen, + nstp, + tsmult, + ) ], - [ - 3, # particleid - 1, # particlegroup - 3, # particleidloc - 2, # status (terminated at boundary face) - 0.0, # time0 - 3.0, # time - 1, # node0 - 1, # k0 - 0.1, # xloc0 - 0.1, # yloc0 - 0.1, # zloc0 - 0.0, # x0 - 1.0, # y0 - 2.0, # z0 - 1, # zone0 - 1, # initialcellface - 5, # node - 2, # k - 0.9, # xloc - 0.9, # yloc - 0.9, # zloc - 10.0, # x - 11.0, # y - 12.0, # z - 2, # zone - 2, # cellface + ) + + # create prt model + prt = flopy.mf6.ModflowPrt(sim, modelname=prt_name, save_flows=True) + + # create prt discretization + flopy.mf6.modflow.mfgwfdis.ModflowGwfdis( + prt, + pname="dis", + nlay=nlay, + nrow=nrow, + ncol=ncol, + ) + + # create mip package + flopy.mf6.ModflowPrtmip( + prt, pname="mip", porosity=porosity + ) + + # create prp package + prp_track_file = f"{prt_name}.prp.trk" + prp_track_csv_file = f"{prt_name}.prp.trk.csv" + flopy.mf6.ModflowPrtprp( + prt, + pname="prp1", + filename=f"{prt_name}_1.prp", + nreleasepts=len(releasepts_prt), + packagedata=releasepts_prt, + perioddata={0: ["FIRST"]}, + track_filerecord=[prp_track_file], + trackcsv_filerecord=[prp_track_csv_file], + stop_at_weak_sink="saws" in prt_name, + boundnames=True, + exit_solve_tolerance=1e-5, + ) + + # create output control package + prt_budget_file = f"{prt_name}.bud" + prt_track_file = f"{prt_name}.trk" + prt_track_csv_file = f"{prt_name}.trk.csv" + flopy.mf6.ModflowPrtoc( + prt, + pname="oc", + budget_filerecord=[prt_budget_file], + track_filerecord=[prt_track_file], + trackcsv_filerecord=[prt_track_csv_file], + saverecord=[("BUDGET", "ALL")], + ) + + # create the flow model interface + gwf_budget_file = gwf_ws / f"{gwf_name}.bud" + gwf_head_file = gwf_ws / f"{gwf_name}.hds" + flopy.mf6.ModflowPrtfmi( + prt, + packagedata=[ + ("GWFHEAD", gwf_head_file), + ("GWFBUDGET", gwf_budget_file), ], - ], - columns=MpEndpointFile.dtypes[7].names, -) + ) + + # add explicit model solution + ems = flopy.mf6.ModflowEms( + sim, + pname="ems", + filename=f"{prt_name}.ems", + ) + sim.register_solution_package(ems, [prt.name]) + + return sim @pytest.mark.parametrize("dataframe", [True, False]) -@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"]) -def test_to_mp7_pathlines(dataframe, source): - if source == "prt": - pls = ( - PRT_TEST_PATHLINES - if dataframe - else PRT_TEST_PATHLINES.to_records(index=False) - ) - elif source == "mp3": - pass - elif source == "mp5": - pass - elif source == "mp7": - pass - mp7_pls = to_mp7_pathlines(pls) +def test_to_mp7_pathlines(prt_sim, dataframe): + prt_pls = pd.read_csv(prt_sim.sim_path / f"{prt_sim.name}.trk.csv") + mp7_pls = to_mp7_pathlines(prt_pls) + assert ( - type(pls) + type(prt_pls) == type(mp7_pls) == (pd.DataFrame if dataframe else np.recarray) ) - assert len(mp7_pls) == 10 assert set( dict(mp7_pls.dtypes).keys() if dataframe else mp7_pls.dtype.names ) == set(MpPathlineFile.dtypes[7].names) @pytest.mark.parametrize("dataframe", [True, False]) -@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"]) -def test_to_mp7_pathlines_empty(dataframe, source): - if source == "prt": - pls = to_mp7_pathlines( - pd.DataFrame.from_records( - [], columns=PrtPathlineFile.dtypes["base"].names - ) - if dataframe - else np.recarray((0,), dtype=PrtPathlineFile.dtypes["base"]) +def test_to_mp7_pathlines_empty(dataframe): + prt_pls = ( + pd.DataFrame.from_records( + [], columns=PrtPathlineFile.dtype.names ) - elif source == "mp3": - pass - elif source == "mp5": - pass - elif source == "mp7": - pass - assert pls.empty if dataframe else pls.size == 0 + if dataframe + else np.recarray((0,), dtype=PrtPathlineFile.dtype) + ) + mp7_pls = to_mp7_pathlines(prt_pls) + + assert prt_pls.empty if dataframe else prt_pls.size == 0 if dataframe: - pls = pls.to_records(index=False) - assert pls.dtype == MpPathlineFile.dtypes[7] + mp7_pls = mp7_pls.to_records(index=False) + assert mp7_pls.dtype == MpPathlineFile.dtypes[7] @pytest.mark.parametrize("dataframe", [True, False]) -def test_to_mp7_pathlines_noop(dataframe): - pls = ( - MP7_TEST_PATHLINES - if dataframe - else MP7_TEST_PATHLINES.to_records(index=False) +def test_to_mp7_pathlines_noop(mp7_sim, dataframe): + plf = flopy.utils.PathlineFile(Path(mp7_sim.model_ws) / f"{mp7_sim.name}.mppth") + og_pls = pd.DataFrame( + plf.get_destination_pathline_data(range(mp7_sim.modelgrid.nnodes), to_recarray=True) ) - mp7_pls = to_mp7_pathlines(pls) + mp7_pls = to_mp7_pathlines(og_pls) + assert ( - type(pls) + type(mp7_pls) == type(mp7_pls) == (pd.DataFrame if dataframe else np.recarray) ) - assert len(mp7_pls) == 2 assert set( dict(mp7_pls.dtypes).keys() if dataframe else mp7_pls.dtype.names ) == set(MpPathlineFile.dtypes[7].names) assert np.array_equal( - mp7_pls if dataframe else pd.DataFrame(mp7_pls), MP7_TEST_PATHLINES + mp7_pls if dataframe else pd.DataFrame(mp7_pls), og_pls ) @pytest.mark.parametrize("dataframe", [True, False]) -@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"]) -def test_to_mp7_endpoints(dataframe, source): - if source == "prt": - eps = to_mp7_endpoints( - PRT_TEST_PATHLINES - if dataframe - else PRT_TEST_PATHLINES.to_records(index=False) - ) - elif source == "mp3": - pass - elif source == "mp5": - pass - elif source == "mp6": - pass - assert len(eps) == 1 - assert np.isclose(eps.time[0], PRT_TEST_PATHLINES.t.max()) +def test_to_mp7_endpoints(prt_sim, dataframe): + prt_pls = pd.read_csv(prt_sim.sim_path / f"{prt_sim.name}.trk.csv") + prt_eps = prt_pls[prt_pls.ireason == 3] + mp7_eps = to_mp7_endpoints(mp7_eps) + + assert np.isclose(mp7_eps.time[0], prt_eps.t.max()) assert set( - dict(eps.dtypes).keys() if dataframe else eps.dtype.names + dict(mp7_eps.dtypes).keys() if dataframe else mp7_eps.dtype.names ) == set(MpEndpointFile.dtypes[7].names) @pytest.mark.parametrize("dataframe", [True, False]) -@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"]) -def test_to_mp7_endpoints_empty(dataframe, source): - eps = to_mp7_endpoints( +def test_to_mp7_endpoints_empty(dataframe): + mp7_eps = to_mp7_endpoints( pd.DataFrame.from_records( - [], columns=PrtPathlineFile.dtypes["base"].names + [], columns=PrtPathlineFile.dtype.names ) if dataframe - else np.recarray((0,), dtype=PrtPathlineFile.dtypes["base"]) + else np.recarray((0,), dtype=PrtPathlineFile.dtype) ) - assert eps.empty if dataframe else eps.size == 0 + + assert mp7_eps.empty if dataframe else mp7_eps.size == 0 if dataframe: - eps = eps.to_records(index=False) - assert eps.dtype == MpEndpointFile.dtypes[7] + mp7_eps = mp7_eps.to_records(index=False) + assert mp7_eps.dtype == MpEndpointFile.dtypes[7] @pytest.mark.parametrize("dataframe", [True, False]) -def test_to_mp7_endpoints_noop(dataframe): +def test_to_mp7_endpoints_noop(mp7_sim, dataframe): """Test a recarray or dataframe which already contains MP7 endpoint data""" - eps = to_mp7_endpoints( - MP7_TEST_ENDPOINTS - if dataframe - else MP7_TEST_ENDPOINTS.to_records(index=False) + epf = flopy.utils.EndpointFile(Path(mp7_sim.model_ws) / f"{mp7_sim.name}.mpend") + og_eps = pd.DataFrame( + epf.get_destination_endpoint_data(range(mp7_sim.modelgrid.nnodes), to_recarray=True) ) + mp7_eps = to_mp7_endpoints(og_eps) + assert np.array_equal( - eps if dataframe else pd.DataFrame(eps), MP7_TEST_ENDPOINTS + mp7_eps if dataframe else pd.DataFrame(mp7_eps), og_eps ) @pytest.mark.parametrize("dataframe", [True, False]) -@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"]) -def test_to_prt_pathlines_roundtrip(dataframe, source): - if source == "prt": - pls = to_mp7_pathlines( - PRT_TEST_PATHLINES - if dataframe - else PRT_TEST_PATHLINES.to_records(index=False) - ) - elif source == "mp3": - pass - elif source == "mp5": - pass - elif source == "mp6": - pass - prt_pls = to_prt_pathlines(pls) +def test_to_prt_pathlines_roundtrip(prt_sim, dataframe): + og_pls = pd.read_csv(prt_sim.sim_path / f"{prt_sim.name}.trk.csv") + mp7_pls = to_mp7_pathlines( + og_pls + if dataframe + else og_pls.to_records(index=False) + ) + prt_pls = to_prt_pathlines(mp7_pls) + if not dataframe: prt_pls = pd.DataFrame(prt_pls) - # import pdb; pdb.set_trace() assert np.allclose( - PRT_TEST_PATHLINES.drop( + prt_pls.drop( ["imdl", "iprp", "irpt", "name", "istatus", "ireason"], axis=1, ), - prt_pls.drop( + og_pls.drop( ["imdl", "iprp", "irpt", "name", "istatus", "ireason"], axis=1, ), @@ -485,25 +395,18 @@ def test_to_prt_pathlines_roundtrip(dataframe, source): @pytest.mark.parametrize("dataframe", [True, False]) -@pytest.mark.parametrize("source", ["prt"]) # , "mp3", "mp5", "mp6"]) -def test_to_prt_pathlines_roundtrip_empty(dataframe, source): - if source == "prt": - pls = to_mp7_pathlines( - pd.DataFrame.from_records( - [], columns=PrtPathlineFile.dtypes["base"].names - ) - if dataframe - else np.recarray((0,), dtype=PrtPathlineFile.dtypes["base"]) +def test_to_prt_pathlines_roundtrip_empty(dataframe): + og_pls = to_mp7_pathlines( + pd.DataFrame.from_records( + [], columns=PrtPathlineFile.dtype.names ) - elif source == "mp3": - pass - elif source == "mp5": - pass - elif source == "mp6": - pass - prt_pls = to_prt_pathlines(pls) - assert pls.empty if dataframe else pls.size == 0 - assert prt_pls.empty if dataframe else pls.size == 0 + if dataframe + else np.recarray((0,), dtype=PrtPathlineFile.dtype) + ) + prt_pls = to_prt_pathlines(og_pls) + + assert og_pls.empty if dataframe else og_pls.size == 0 + assert prt_pls.empty if dataframe else og_pls.size == 0 assert set( - dict(pls.dtypes).keys() if dataframe else pls.dtype.names + dict(og_pls.dtypes).keys() if dataframe else og_pls.dtype.names ) == set(MpPathlineFile.dtypes[7].names) diff --git a/autotest/test_prtfile.py b/autotest/test_prtfile.py index 39b7850fc..61c5dfda2 100644 --- a/autotest/test_prtfile.py +++ b/autotest/test_prtfile.py @@ -1,45 +1,60 @@ -from pathlib import Path - import pytest from autotest.conftest import get_project_root_path from flopy.utils.prtfile import PathlineFile pytestmark = pytest.mark.mf6 +proj_root = get_project_root_path() prt_data_path = ( - get_project_root_path() / "examples" / "data" / "mf6" / "prt_data" / "001" + proj_root / "examples" / "data" / "mf6" / "prt_data" / "001" ) @pytest.mark.parametrize( "path, header_path", [ - (Path(prt_data_path) / "prt001.trk", None), ( - Path(prt_data_path) / "prt001.trk", - Path(prt_data_path) / "prt001.trk.hdr", + prt_data_path / "prt001.trk", + prt_data_path / "prt001.trk.hdr", ), - (Path(prt_data_path) / "prt001.trk.csv", None), + (prt_data_path / "prt001.trk.csv", None), ], ) def test_init(path, header_path): file = PathlineFile(path, header_filename=header_path) assert file.fname == path - assert file.dtype == PathlineFile.dtypes["full"] if path.suffix == ".csv": assert len(file._data) == len(open(path).readlines()) - 1 @pytest.mark.parametrize( - "path", + "path, header_path", [ - Path(prt_data_path) / "prt001.trk", - Path(prt_data_path) / "prt001.trk.csv", + ( + prt_data_path / "prt001.trk", + prt_data_path / "prt001.trk.hdr", + ), + (prt_data_path / "prt001.trk.csv", None), ], ) -def test_intersect(path): - file = PathlineFile(path) +def test_intersect(path, header_path): + file = PathlineFile(path, header_filename=header_path) nodes = [1, 11, 21] intersection = file.intersect(nodes) assert any(intersection) assert all(d.icell in nodes for d in intersection.itertuples()) + + +@pytest.mark.parametrize( + "path, header_path", + [ + ( + prt_data_path / "prt001.trk", + prt_data_path / "prt001.trk.hdr", + ), + (prt_data_path / "prt001.trk.csv", None), + ], +) +def test_validate(path, header_path): + file = PathlineFile(path, header_filename=header_path) + file.validate() diff --git a/flopy/plot/plotutil.py b/flopy/plot/plotutil.py index 320c1ae76..fba2dc062 100644 --- a/flopy/plot/plotutil.py +++ b/flopy/plot/plotutil.py @@ -2733,7 +2733,7 @@ def to_mp7_pathlines( ], dtype=MpPathlineFile.dtypes[7], ) - elif all(n in data.dtypes for n in ParticleTrackFile.dtype.names): + elif all(n in data.dtypes for n in ParticleTrackFile._dtype.names): data = np.core.records.fromarrays( [ data["particleid"], @@ -2762,7 +2762,7 @@ def to_mp7_pathlines( f"{pformat(MpPathlineFile.dtypes[6].names)} for MODPATH 6, or;\n" f"{pformat(MpPathlineFile.dtypes[7].names)} for MODPATH 7, or;\n" f"{pformat(PrtPathlineFile.dtypes['base'].names)} for MODFLOW 6 PRT, or;\n" - f"{pformat(ParticleTrackFile.dtype.names)} (minimal expected dtype names)" + f"{pformat(ParticleTrackFile._dtype.names)} (minimal expected dtype names)" ) return pd.DataFrame(data) if rt == pd.DataFrame else data @@ -2893,7 +2893,7 @@ def to_mp7_endpoints( ], dtype=MpEndpointFile.dtypes[7], ) - elif all(n in data.dtypes for n in ParticleTrackFile.dtype.names): + elif all(n in data.dtypes for n in ParticleTrackFile._dtype.names): data = np.core.records.fromarrays( [ endpts["particleid"], @@ -2933,7 +2933,7 @@ def to_mp7_endpoints( f"{pformat(MpEndpointFile.dtypes[6].names)} for MODPATH 6, or;\n" f"{pformat(MpEndpointFile.dtypes[7].names)} for MODPATH 7, or;\n" f"{pformat(PrtPathlineFile.dtypes['base'].names)} for MODFLOW 6 PRT, or;\n" - f"{pformat(ParticleTrackFile.dtype.names)} (minimal expected dtype names)" + f"{pformat(ParticleTrackFile._dtype.names)} (minimal expected dtype names)" ) return pd.DataFrame(data) if rt == pd.DataFrame else data @@ -3020,7 +3020,7 @@ def to_prt_pathlines( n in data.dtypes for n in PrtPathlineFile.dtypes["base"].names ): # already in prt format? return data if rt == pd.DataFrame else data.to_records(index=False) - elif all(n in data.dtypes for n in ParticleTrackFile.dtype.names): + elif all(n in data.dtypes for n in ParticleTrackFile._dtype.names): data = np.core.records.fromarrays( [ np.zeros(data.shape[0]), @@ -3053,7 +3053,7 @@ def to_prt_pathlines( f"{pformat(MpEndpointFile.dtypes[5].names)} for MODPATH 5 endpoints, or;\n" f"{pformat(MpEndpointFile.dtypes[3].names)} for MODPATH 3 endpoints, or;\n" f"{pformat(PrtPathlineFile.dtypes['base'].names)} for MODFLOW 6 PRT, or;\n" - f"{pformat(ParticleTrackFile.dtype.names)} (minimal expected dtype names)" + f"{pformat(ParticleTrackFile._dtype.names)} (minimal expected dtype names)" ) if rt == pd.DataFrame: diff --git a/flopy/utils/modpathfile.py b/flopy/utils/modpathfile.py index 4fe026fb5..a226d89f3 100644 --- a/flopy/utils/modpathfile.py +++ b/flopy/utils/modpathfile.py @@ -4,6 +4,7 @@ import itertools import os +from abc import ABC, abstractmethod from typing import List, Optional, Tuple, Union import numpy as np @@ -14,9 +15,14 @@ from ..utils.flopy_io import loadtxt -class ModpathFile(ParticleTrackFile): +class ModpathFile(ParticleTrackFile, ABC): """Provides MODPATH output file support.""" + @property + @abstractmethod + def dtypes(self): + pass + def __init__( self, filename: Union[str, os.PathLike], verbose: bool = False ): @@ -161,64 +167,70 @@ class PathlineFile(ModpathFile): """ - dtypes = { - **dict.fromkeys( - [3, 5], - np.dtype( + @property + def dtypes(self): + return { + **dict.fromkeys( + [3, 5], + np.dtype( + [ + ("particleid", np.int32), + ("x", np.float32), + ("y", np.float32), + ("zloc", np.float32), + ("z", np.float32), + ("time", np.float32), + ("j", np.int32), + ("i", np.int32), + ("k", np.int32), + ("cumulativetimestep", np.int32), + ] + ), + ), + 6: np.dtype( [ ("particleid", np.int32), + ("particlegroup", np.int32), + ("timepointindex", np.int32), + ("cumulativetimestep", np.int32), + ("time", np.float32), ("x", np.float32), ("y", np.float32), - ("zloc", np.float32), ("z", np.float32), - ("time", np.float32), - ("j", np.int32), + ("k", np.int32), ("i", np.int32), + ("j", np.int32), + ("grid", np.int32), + ("xloc", np.float32), + ("yloc", np.float32), + ("zloc", np.float32), + ("linesegmentindex", np.int32), + ] + ), + 7: np.dtype( + [ + ("particleid", np.int32), + ("particlegroup", np.int32), + ("sequencenumber", np.int32), + ("particleidloc", np.int32), + ("time", np.float32), + ("x", np.float32), + ("y", np.float32), + ("z", np.float32), ("k", np.int32), - ("cumulativetimestep", np.int32), + ("node", np.int32), + ("xloc", np.float32), + ("yloc", np.float32), + ("zloc", np.float32), + ("stressperiod", np.int32), + ("timestep", np.int32), ] ), - ), - 6: np.dtype( - [ - ("particleid", np.int32), - ("particlegroup", np.int32), - ("timepointindex", np.int32), - ("cumulativetimestep", np.int32), - ("time", np.float32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("k", np.int32), - ("i", np.int32), - ("j", np.int32), - ("grid", np.int32), - ("xloc", np.float32), - ("yloc", np.float32), - ("zloc", np.float32), - ("linesegmentindex", np.int32), - ] - ), - 7: np.dtype( - [ - ("particleid", np.int32), - ("particlegroup", np.int32), - ("sequencenumber", np.int32), - ("particleidloc", np.int32), - ("time", np.float32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("k", np.int32), - ("node", np.int32), - ("xloc", np.float32), - ("yloc", np.float32), - ("zloc", np.float32), - ("stressperiod", np.int32), - ("timestep", np.int32), - ] - ), - } + } + + @property + def dtype(self): + return self._dtype kijnames = [ "k", @@ -236,7 +248,7 @@ def __init__( self, filename: Union[str, os.PathLike], verbose: bool = False ): super().__init__(filename, verbose=verbose) - self.dtype, self._data = self._load() + self._dtype, self._data = self._load() self.nid = np.unique(self._data["particleid"]) def _load(self) -> Tuple[np.dtype, np.ndarray]: @@ -447,98 +459,104 @@ class EndpointFile(ModpathFile): """ - dtypes = { - **dict.fromkeys( - [3, 5], - np.dtype( + @property + def dtypes(self): + return { + **dict.fromkeys( + [3, 5], + np.dtype( + [ + ("zone", np.int32), + ("j", np.int32), + ("i", np.int32), + ("k", np.int32), + ("x", np.float32), + ("y", np.float32), + ("z", np.float32), + ("zloc", np.float32), + ("time", np.float32), + ("x0", np.float32), + ("y0", np.float32), + ("zloc0", np.float32), + ("j0", np.int32), + ("i0", np.int32), + ("k0", np.int32), + ("zone0", np.int32), + ("cumulativetimestep", np.int32), + ("ipcode", np.int32), + ("time0", np.float32), + ] + ), + ), + 6: np.dtype( [ - ("zone", np.int32), - ("j", np.int32), - ("i", np.int32), + ("particleid", np.int32), + ("particlegroup", np.int32), + ("status", np.int32), + ("time0", np.float32), + ("time", np.float32), + ("initialgrid", np.int32), + ("k0", np.int32), + ("i0", np.int32), + ("j0", np.int32), + ("cellface0", np.int32), + ("zone0", np.int32), + ("xloc0", np.float32), + ("yloc0", np.float32), + ("zloc0", np.float32), + ("x0", np.float32), + ("y0", np.float32), + ("z0", np.float32), + ("finalgrid", np.int32), ("k", np.int32), + ("i", np.int32), + ("j", np.int32), + ("cellface", np.int32), + ("zone", np.int32), + ("xloc", np.float32), + ("yloc", np.float32), + ("zloc", np.float32), ("x", np.float32), ("y", np.float32), ("z", np.float32), - ("zloc", np.float32), + ("label", "|S40"), + ] + ), + 7: np.dtype( + [ + ("particleid", np.int32), + ("particlegroup", np.int32), + ("particleidloc", np.int32), + ("status", np.int32), + ("time0", np.float32), ("time", np.float32), + ("node0", np.int32), + ("k0", np.int32), + ("xloc0", np.float32), + ("yloc0", np.float32), + ("zloc0", np.float32), ("x0", np.float32), ("y0", np.float32), - ("zloc0", np.float32), - ("j0", np.int32), - ("i0", np.int32), - ("k0", np.int32), + ("z0", np.float32), ("zone0", np.int32), - ("cumulativetimestep", np.int32), - ("ipcode", np.int32), - ("time0", np.float32), + ("initialcellface", np.int32), + ("node", np.int32), + ("k", np.int32), + ("xloc", np.float32), + ("yloc", np.float32), + ("zloc", np.float32), + ("x", np.float32), + ("y", np.float32), + ("z", np.float32), + ("zone", np.int32), + ("cellface", np.int32), ] ), - ), - 6: np.dtype( - [ - ("particleid", np.int32), - ("particlegroup", np.int32), - ("status", np.int32), - ("time0", np.float32), - ("time", np.float32), - ("initialgrid", np.int32), - ("k0", np.int32), - ("i0", np.int32), - ("j0", np.int32), - ("cellface0", np.int32), - ("zone0", np.int32), - ("xloc0", np.float32), - ("yloc0", np.float32), - ("zloc0", np.float32), - ("x0", np.float32), - ("y0", np.float32), - ("z0", np.float32), - ("finalgrid", np.int32), - ("k", np.int32), - ("i", np.int32), - ("j", np.int32), - ("cellface", np.int32), - ("zone", np.int32), - ("xloc", np.float32), - ("yloc", np.float32), - ("zloc", np.float32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("label", "|S40"), - ] - ), - 7: np.dtype( - [ - ("particleid", np.int32), - ("particlegroup", np.int32), - ("particleidloc", np.int32), - ("status", np.int32), - ("time0", np.float32), - ("time", np.float32), - ("node0", np.int32), - ("k0", np.int32), - ("xloc0", np.float32), - ("yloc0", np.float32), - ("zloc0", np.float32), - ("x0", np.float32), - ("y0", np.float32), - ("z0", np.float32), - ("zone0", np.int32), - ("initialcellface", np.int32), - ("node", np.int32), - ("k", np.int32), - ("xloc", np.float32), - ("yloc", np.float32), - ("zloc", np.float32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("zone", np.int32), - ("cellface", np.int32), - ] - ), - } + } + + @property + def dtype(self): + return self._dtype kijnames = [ "k0", @@ -560,7 +578,7 @@ def __init__( self, filename: Union[str, os.PathLike], verbose: bool = False ): super().__init__(filename, verbose) - self.dtype, self._data = self._load() + self._dtype, self._data = self._load() self.nid = np.unique(self._data["particleid"]) def _load(self) -> Tuple[np.dtype, np.ndarray]: @@ -799,61 +817,67 @@ class TimeseriesFile(ModpathFile): >>> ts1 = tsobj.get_data(partid=1) """ - dtypes = { - **dict.fromkeys( - [3, 5], - np.dtype( + @property + def dtypes(self): + return { + **dict.fromkeys( + [3, 5], + np.dtype( + [ + ("timestepindex", np.int32), + ("particleid", np.int32), + ("node", np.int32), + ("x", np.float32), + ("y", np.float32), + ("z", np.float32), + ("zloc", np.float32), + ("time", np.float32), + ("timestep", np.int32), + ] + ), + ), + 6: np.dtype( [ - ("timestepindex", np.int32), + ("timepointindex", np.int32), + ("timestep", np.int32), + ("time", np.float32), ("particleid", np.int32), - ("node", np.int32), + ("particlegroup", np.int32), ("x", np.float32), ("y", np.float32), ("z", np.float32), + ("grid", np.int32), + ("k", np.int32), + ("i", np.int32), + ("j", np.int32), + ("xloc", np.float32), + ("yloc", np.float32), ("zloc", np.float32), - ("time", np.float32), + ] + ), + 7: np.dtype( + [ + ("timepointindex", np.int32), ("timestep", np.int32), + ("time", np.float32), + ("particleid", np.int32), + ("particlegroup", np.int32), + ("particleidloc", np.int32), + ("node", np.int32), + ("xloc", np.float32), + ("yloc", np.float32), + ("zloc", np.float32), + ("x", np.float32), + ("y", np.float32), + ("z", np.float32), + ("k", np.int32), ] ), - ), - 6: np.dtype( - [ - ("timepointindex", np.int32), - ("timestep", np.int32), - ("time", np.float32), - ("particleid", np.int32), - ("particlegroup", np.int32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("grid", np.int32), - ("k", np.int32), - ("i", np.int32), - ("j", np.int32), - ("xloc", np.float32), - ("yloc", np.float32), - ("zloc", np.float32), - ] - ), - 7: np.dtype( - [ - ("timepointindex", np.int32), - ("timestep", np.int32), - ("time", np.float32), - ("particleid", np.int32), - ("particlegroup", np.int32), - ("particleidloc", np.int32), - ("node", np.int32), - ("xloc", np.float32), - ("yloc", np.float32), - ("zloc", np.float32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("k", np.int32), - ] - ), - } + } + + @property + def dtype(self): + return self._dtype kijnames = [ "k", @@ -870,7 +894,7 @@ class TimeseriesFile(ModpathFile): def __init__(self, filename, verbose=False): super().__init__(filename, verbose) - self.dtype, self._data = self._load() + self._dtype, self._data = self._load() self.nid = np.unique(self._data["particleid"]) def _load(self) -> Tuple[np.dtype, np.ndarray]: diff --git a/flopy/utils/particletrackfile.py b/flopy/utils/particletrackfile.py index bfd481178..186267fb7 100644 --- a/flopy/utils/particletrackfile.py +++ b/flopy/utils/particletrackfile.py @@ -4,11 +4,13 @@ import os from abc import ABC, abstractmethod +from collections import OrderedDict from pathlib import Path from typing import Union import numpy as np from numpy.lib.recfunctions import stack_arrays +import pandas as pd MIN_PARTICLE_TRACK_DTYPE = np.dtype( [ @@ -27,10 +29,6 @@ class ParticleTrackFile(ABC): Abstract base class for particle track output files. Exposes a unified API supporting MODPATH versions 3, 5, 6 and 7, as well as MODFLOW 6 PRT models. - Notes - ----- - - Parameters ---------- filename : str or PathLike @@ -40,17 +38,20 @@ class ParticleTrackFile(ABC): """ - dtype = MIN_PARTICLE_TRACK_DTYPE - """ - Minimal data shared by all particle track file formats. - """ - # legacy, todo: deprecate outdtype = MIN_PARTICLE_TRACK_DTYPE """ Minimal data shared by all particle track file formats. """ + @property + @abstractmethod + def dtype(self): + """ + Particle track file data dtype. + """ + return MIN_PARTICLE_TRACK_DTYPE + def __init__( self, filename: Union[str, os.PathLike], @@ -334,3 +335,16 @@ def write_shapefile( # write the final recarray to a shapefile recarray2shp(sdata, geoms, shpname=shpname, crs=crs, **kwargs) + + def validate(self): + dtype = self.dtype + expected = OrderedDict(MIN_PARTICLE_TRACK_DTYPE.fields.items()) + if isinstance(dtype, dict): + for dt in dtype.values(): + self.validate(dt) + elif isinstance(dtype, pd.Series): + subset = OrderedDict({k: v for k, v in dtype.to_dict().items() if k in MIN_PARTICLE_TRACK_DTYPE.names}) + assert subset == expected + elif isinstance(dtype, np.dtypes.VoidDType): + subset = OrderedDict({k: v for k, v in dtype.fields.items() if k in MIN_PARTICLE_TRACK_DTYPE.names}) + assert subset == expected \ No newline at end of file diff --git a/flopy/utils/prtfile.py b/flopy/utils/prtfile.py index 6df241974..219ef6d3b 100644 --- a/flopy/utils/prtfile.py +++ b/flopy/utils/prtfile.py @@ -18,61 +18,18 @@ class PathlineFile(ParticleTrackFile): key_cols = ["imdl", "iprp", "irpt", "trelease"] """Columns making up each particle's composite key.""" - dtypes = { - "base": np.dtype( - [ - ("kper", np.int32), - ("kstp", np.int32), - ("imdl", np.int32), - ("iprp", np.int32), - ("irpt", np.int32), - ("ilay", np.int32), - ("icell", np.int32), - ("izone", np.int32), - ("istatus", np.int32), - ("ireason", np.int32), - ("trelease", np.float32), - ("t", np.float32), - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("name", np.str_), - ] - ), - "full": np.dtype( - [ - ("kper", np.int32), - ("kstp", np.int32), - ("imdl", np.int32), - ("iprp", np.int32), - ("irpt", np.int32), - ("ilay", np.int32), - ("k", np.int32), # conform to canonical dtype - ("icell", np.int32), - ("izone", np.int32), - ("idest", np.int32), # destination zone, for convenience - ("dest", np.str_), # destination name, for convenience - ("istatus", np.int32), - ("ireason", np.int32), - ("trelease", np.float32), - ("t", np.float32), - ("t0", np.float32), # release time, for convenience - ("tt", np.float32), # termination time, convenience - ("time", np.float32), # conform to canonical dtype - ("x", np.float32), - ("y", np.float32), - ("z", np.float32), - ("particleid", np.int32), # conform to canonical dtype - ("name", np.str_), - ] - ), - } - """ - Base (directly from MODFLOW 6 PRT) and full (extended by this class) dtypes for PRT. - The extended dtype complies with the canonical `flopy.utils.particletrackfile.dtype` - and provides a few other columns for convenience, including release and termination - time and destination zone number/name. - """ + @property + def dtype(self): + """ + PRT track file dtype. This is loaded dynamically from the binary header file or CSV file + headers. A best-effort attempt is made to add extra columns to comply with the canonical + `flopy.utils.particletrackfile.dtype`, as well as convenience columns, including release + and termination time and destination zone number and name. + + Consumers of this class which expect the canonical particle track file attributes should + call `validate()` to make sure the attributes were successfully loaded. + """ + return self._dtype def __init__( self, @@ -82,14 +39,14 @@ def __init__( verbose: bool = False, ): super().__init__(filename, verbose) - self.dtype, self._data = self._load(header_filename, destination_map) + self._dtype, self._data = self._load(header_filename, destination_map) def _load( self, header_filename=None, destination_map=None, ) -> np.ndarray: - # load dtype from header file if present, otherwise use default dtype + # if a header file is present, we're reading a binary file hdr_fname = ( None if header_filename is None @@ -103,17 +60,11 @@ def _load( "formats": lines[1].strip().split(","), } ) + data = pd.DataFrame(np.fromfile(self.fname, dtype=dtype)) else: - dtype = self.dtypes["base"] - - # read as binary or csv - try: - data = pd.read_csv(self.fname, dtype=dtype) - except UnicodeDecodeError: - try: - data = pd.DataFrame(np.fromfile(self.fname, dtype=dtype)) - except: - raise ValueError("Unreadable file, expected binary or CSV") + # otherwise we're reading a CSV file + data = pd.read_csv(self.fname) + dtype = data.to_records(index=False).dtype # add particle id column data = data.sort_values(self.key_cols) @@ -134,6 +85,9 @@ def _load( .tt ) + # add k column + data["k"] = data["ilay"] + # assign destinations if zone map is provided if destination_map is not None: data["idest"] = data[data.istatus > 1].izone @@ -142,7 +96,7 @@ def _load( axis=1, ) - return self.dtypes["full"], data + return data.to_records(index=False).dtype, data def intersect( self, cells, to_recarray=True @@ -160,3 +114,12 @@ def intersect( return sect[idxs].sort_values(by=["particleid", "time"]) else: return [self.get_data(pid) for pid in pids] + + @staticmethod + def get_track_dtype(path: Union[str, os.PathLike]): + """Read a numpy dtype describing particle track + data format from the ascii track header file.""" + + hdr_lns = open(path).readlines() + hdr_lns_spl = [[ll.strip() for ll in l.split(",")] for l in hdr_lns] + return np.dtype(list(zip(hdr_lns_spl[0], hdr_lns_spl[1])))