def animate(t):
global pf
n, s = divmod(t, numrec)
if s == 0:
pf = ParticleFile(particle_files[n])
X, Y = pf.position(s)
particle_dist.set_data(X, Y)
timestamp.set_text(pf.time(s))
return particle_dist, timestamp
def test_particle_variable(particle_file):
"""Two particle variables, start_time and location_id"""
with ParticleFile(particle_file) as pf:
assert pf.start_time[0] == np.datetime64("1970-01-01")
assert pf["start_time"][1] == np.datetime64("1970-01-01 01")
assert all(pf.location_id == np.array([10000, 10001, 10002]))
assert all(pf.location_id == pf["location_id"][:])
def test_position(particle_file):
with ParticleFile(particle_file) as pf:
X, Y = pf.position(time=1)
assert all(X == pf.X[1])
assert all(Y == pf.Y[1])
X, Y = pf.position(2)
assert all(X == pf.X[2])
assert all(Y == pf.Y[2])
def test_count(particle_file):
"""Alignment of time frames in the particle file."""
with ParticleFile(particle_file) as pf:
assert pf.num_times == 4
assert all(pf.start == [0, 1, 3, 5])
assert list(pf.count) == [1, 2, 2, 1]
assert list(pf.end) == [1, 3, 5, 6]
assert len(pf) == 4
assert pf.num_particles == 3
def test_time(particle_file):
"""Time handled correctly"""
with ParticleFile(particle_file) as pf:
assert pf.time[3] == np.datetime64("1970-01-01 03")
times2 = [np.datetime64(t) for t in ["1970-01-01", "1970-01-01 01"]]
assert all(pf.time[:2] == times2)
# Old callable notation still works
assert pf.time(3) == pf.time[3]
assert str(pf.time(3)) == "1970-01-01T03:00:00"
def test_pid(particle_file):
"""The pid is correct"""
with ParticleFile(particle_file) as pf:
assert pf.pid.isel(time=0) == 0
assert pf["pid"][0] == 0
assert pf.pid[0] == 0
assert all(pf.pid[1] == [0, 1])
assert list(pf.pid[2]) == [0, 2]
assert pf.pid[3] == 2
def test_pid2(particle_file):
"""The pid from an instance variable"""
with ParticleFile(particle_file) as pf:
X = pf.X
assert all(X.pid == pf.ds.pid)
for i in range(4):
assert all(X[i].pid == pf.pid[i])
# X.pid[3] is not the same as X[3].pid
assert not all(X.pid[3] == X[3].pid)
def test_isel(particle_file):
with ParticleFile(particle_file) as pf:
X = pf.X
assert all(X.isel(time=2) == X[2])
with pytest.raises(TypeError):
# Need an argument
X.isel()
with pytest.raises(TypeError):
# Only keyword arguments
X.isel(2)
def test_trajectory(particle_file):
with ParticleFile(particle_file) as pf:
X, Y = pf.trajectory(2)
assert all(X == [22, 23])
assert all(Y == [9, 10])
traj = pf.trajectory(0)
assert len(traj) == 3
assert all(traj.time == pf.time[:-1])
assert all(traj.X == pf.X.sel(pid=0))
assert all(traj.Y == pf.Y.sel(pid=0))
def test_getX(particle_file):
with ParticleFile(particle_file) as pf:
X = pf.X
assert X == pf["X"]
assert X == pf.variables["X"] # Obsolete
assert X.isel(time=0) == 0
assert X[0] == 0
assert X[0] == 0
assert all(X[1] == [1, 11])
assert all(X[2] == [2, 22])
assert X[3] == 23
def test_X_slice(particle_file):
"""Can read variables with time slices"""
with ParticleFile(particle_file) as pf:
X = pf.X
V = pf.X[1:3]
assert len(V) == 2
assert all(V[0] == X[1])
assert all(V[1] == X[2])
assert all(V.da == [1, 11, 2, 22])
assert all(V.count == X.count[1:3])
assert all(V.time == X.time[1:3])
assert all(V.pid == [0, 1, 0, 2])
V = X[:]
assert (V.da == X.da).all()
with pytest.raises(IndexError):
pf.X[::2] # Do not accept strides != 1
def test_full(particle_file):
with ParticleFile(particle_file) as pf:
X = pf.X
V = X.full()
assert V[0, 0] == 0
assert np.isnan(V[0, 1])
assert np.isnan(V[0, 2])
assert V[1, 0] == 1
assert V[1, 1] == 11
assert np.isnan(V[1, 2])
assert V[2, 0] == 2
assert np.isnan(V[2, 1])
assert V[2, 2] == 22
assert np.isnan(V[3, 0])
assert np.isnan(V[3, 1])
assert V[3, 2] == 23
def test_sel(particle_file):
with ParticleFile(particle_file) as pf:
X = pf.X
assert all(X.sel(pid=0) == [0, 1, 2])
assert all(X.sel(pid=1) == [11])
assert all(X.sel(pid=2) == [22, 23])
assert all(X.sel(time=X.time[2]) == X[2])
assert all(X.sel(time="1970-01-01 02") == X[2])
assert all(X.sel(time=np.datetime64("1970-01-01 02")) == X[2])
assert all(X.sel(time=datetime.datetime(1970, 1, 1, 2)) == X[2])
assert X.sel(time="1970-01-01 02", pid=0) == X[2, 0]
assert X.sel(pid=2, time="1970-01-01 03") == X[3, 2]
with pytest.raises(KeyError):
X.sel(time="1980-01-01")
with pytest.raises(KeyError):
X.sel(pid=-1)
with pytest.raises(ValueError):
X.sel()
with pytest.raises(TypeError):
X.sel("1970-01-01 02")
# ROMS grid, plot domain
with Dataset(grid_file) as f0:
H = f0.variables["h"][j0:j1, i0:i1]
M = f0.variables["mask_rho"][j0:j1, i0:i1]
lon = f0.variables["lon_rho"][j0:j1, i0:i1]
lat = f0.variables["lat_rho"][j0:j1, i0:i1]
# Cell centers and boundaries
Xcell = np.arange(i0, i1)
Ycell = np.arange(j0, j1)
Xb = np.arange(i0 - 0.5, i1)
Yb = np.arange(j0 - 0.5, j1)
# particle_file
pf = ParticleFile(particle_file)
num_times = pf.num_times
# Set up the plot area
fig = plt.figure(figsize=(12, 10))
ax = plt.axes(xlim=(i0 + 1, i1 - 1), ylim=(j0 + 1, j1 - 1), aspect="equal")
# Background bathymetry
cmap = plt.get_cmap("Blues")
ax.contourf(Xcell, Ycell, H, cmap=cmap, alpha=0.3)
# Lon/lat lines
ax.contour(Xcell,
Ycell,
lat,
levels=range(54, 61),
# ---------------
# Files
particle_file = "latlon.nc"
# time step to plot
t = 90
# Geographical extent
lon0, lat0, lon1, lat1 = -6, 54, 12, 62
# --------------------
# Read particle_file
# --------------------
with ParticleFile(particle_file) as pf:
lon = pf["lon"][t]
lat = pf["lat"][t]
# -------------------------
# Make background map
# -------------------------
# Polarstereographic projection, northern North Sea
m = Basemap(
projection="stere",
resolution="i",
llcrnrlon=lon0,
llcrnrlat=lat0,
urcrnrlon=lon1,
urcrnrlat=lat1,
# ROMS grid, plot domain
with Dataset(grid_file) as f0:
H = f0.variables["h"][j0:j1, i0:i1]
M = f0.variables["mask_rho"][j0:j1, i0:i1]
lon = f0.variables["lon_rho"][j0:j1, i0:i1]
lat = f0.variables["lat_rho"][j0:j1, i0:i1]
# Cell centers and boundaries
Xcell = np.arange(i0, i1)
Ycell = np.arange(j0, j1)
Xb = np.arange(i0 - 0.5, i1)
Yb = np.arange(j0 - 0.5, j1)
# particle_file
pf = ParticleFile(particle_file)
num_times = pf.num_times
# Set up the plot area
fig = plt.figure(figsize=(12, 10))
ax = plt.axes(xlim=(i0 + 1, i1 - 1), ylim=(j0 + 1, j1 - 1), aspect="equal")
# Background bathymetry
cmap = plt.get_cmap("Blues")
ax.contourf(Xcell, Ycell, H, cmap=cmap, alpha=0.3)
# Lon/lat lines
ax.contour(Xcell, Ycell, lat, levels=range(56, 83, 2), colors="black", linestyles=":")
ax.contour(Xcell, Ycell, lon, levels=range(-4, 40, 5), colors="black", linestyles=":")
# Landmask
# ROMS grid, plot domain
with Dataset(grid_file) as f0:
H = f0.variables["h"][j0:j1, i0:i1]
M = f0.variables["mask_rho"][j0:j1, i0:i1]
lon = f0.variables["lon_rho"][j0:j1, i0:i1]
lat = f0.variables["lat_rho"][j0:j1, i0:i1]
# Cell centers and boundaries
Xcell = np.arange(i0, i1)
Ycell = np.arange(j0, j1)
Xb = np.arange(i0 - 0.5, i1)
Yb = np.arange(j0 - 0.5, j1)
# particle_file
pf = ParticleFile(particle_file)
# Make plot
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(1, 1, 1)
# Background map
# Bathymetry
cmap = plt.get_cmap("Blues")
h = ax.contourf(Xcell, Ycell, H, cmap=cmap, alpha=0.9)
# Landmask
constmap = plt.matplotlib.colors.ListedColormap([0.2, 0.6, 0.4])
M = np.ma.masked_where(M > 0, M)
plt.pcolormesh(Xb, Yb, M, cmap=constmap)
i0, i1 = 100, 130
j0, j1 = 90, 115
# timestamp
t = 176
# ----------------
# ROMS grid, plot domain
# Slight overkill to use roppy, could be more stand alone
f0 = Dataset(grid_file)
g = roppy.SGrid(f0, subgrid=(i0, i1, j0, j1))
# particle_file
pf = ParticleFile(particle_file)
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(1, 1, 1)
# Make background map
cmap = plt.get_cmap("Blues")
h = ax.contourf(g.X, g.Y, g.h, cmap=cmap, alpha=0.3)
roppy.mpl_util.landmask(g, color=(0.6, 0.8, 0.0))
ax.contour(g.X, g.Y, g.lat_rho, levels=range(57, 64), colors="black", linestyles=":")
ax.contour(g.X, g.Y, g.lon_rho, levels=range(-4, 10, 2), colors="black", linestyles=":")
# Plot particle distribution
X, Y = pf.position(time=t)
timestring = pf.time(t)
tstep = 40 # Time step to show
# Output file (and type)
outfile = "line_bokeh.png"
#outfile = "line_bokeh.html"
scale = 3 # Figure size factor
# --- Data files ---
ladim_dir = Path("../../")
grid_file = ladim_dir / "examples/data/ocean_avg_0014.nc"
particle_file = ladim_dir / "examples/line/line.nc"
# --- Read particle data ---
pf = ParticleFile(particle_file)
X = pf.X[tstep].values
Y = pf.Y[tstep].values
# --- Background bathymetry data ---
# Read bathymetry and land mask
with Dataset(grid_file) as f:
H = f.variables["h"][:, :]
M = f.variables["mask_rho"][:, :]
jmax, imax = M.shape
H[M < 1] = np.nan # Mask out land
# --- Plot ---
# Figure
restarted_file = "restart_0002.nc"
grid_file = "../data/ocean_avg_0014.nc"
# Subgrid definition
i0, i1 = 113, 119
j0, j1 = 97, 102
# Comparable records,split by 4
t0 = 9 # tenth record in unsplit file
t1 = 1 # second record in third split file
t2 = 1 # second record in first restarted file
# ----------------
# particle_file
pf0 = ParticleFile(unsplit_file)
pf1 = ParticleFile(split_file)
pf2 = ParticleFile(restarted_file)
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(1, 1, 1)
# Make bathymetry background
with Dataset(grid_file) as f0:
H = f0.variables["h"][j0:j1, i0:i1]
cmap = plt.get_cmap("Blues")
ax.contourf(range(i0, i1), range(j0, j1), H, cmap=cmap, alpha=0.3)
# Plot initial particle distribution
X0, Y0 = pf0.position(t0)
X1, Y1 = pf1.position(t1)
def test_values(particle_file):
with ParticleFile(particle_file) as pf:
X = pf.X
assert all(X.values == X.da.values)
# --- Settings ---
tstep = 40 # Time step to show
# Output file (and type)
output_file = "line_hv.png"
#output_file = "line_hv.html"
scale = 5 # Figure size factor
# --- Data files ---
ladim_dir = Path("../../")
grid_file = ladim_dir / "examples/data/ocean_avg_0014.nc"
particle_file = ladim_dir / "examples/line/line.nc"
# --- Read particle data ---
pf = ParticleFile(particle_file)
X, Y = pf.position(tstep)
# --- Background bathymetry data ---
# Read bathymetry and land mask
with xr.open_dataset(grid_file) as A:
H = A.h
M = A.mask_rho
jmax, imax = M.shape
H = H.where(M > 0) # Mask out land
M = M.where(M < 1) # Mask out sea
# --- Holoviews elements ---
# Land image
import numpy as np
import matplotlib.pyplot as plt
from postladim import ParticleFile, cellcount
pf = ParticleFile("line.nc")
X, Y = pf.position(80)
C = cellcount(X, Y)
i0 = int(round(X.min()))
j0 = int(round(Y.min()))
jmax, imax = C.shape
x_edges = np.arange(i0 - 0.5, i0 + imax)
y_edges = np.arange(j0 - 0.5, j0 + jmax)
plt.set_cmap("magma_r")
plt.pcolormesh(x_edges, y_edges, C)
plt.colorbar()
plt.plot(X, Y, ".k")
plt.show()
# import itertools import numpy as np import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation from postladim import ParticleFile from gridforce_discrete import Grid # Input file particle_file = "obstacle.nc" # Get the grid information grid = Grid(None) # Open the particle_file pf = ParticleFile(particle_file) num_times = pf.num_times # Set up the plot area fig = plt.figure(figsize=(12, 8)) ax = plt.axes(xlim=(0, grid.imax), ylim=(0, grid.jmax), aspect="equal") # Landmask Xb = np.arange(-0.5, grid.imax) Yb = np.arange(-0.5, grid.jmax) constmap = plt.matplotlib.colors.ListedColormap([0.2, 0.6, 0.4]) M = np.ma.masked_where(grid.M > 0, grid.M) plt.pcolormesh(Xb, Yb, M, cmap=constmap) # Draw the cirular boundary T = np.linspace(0, np.pi) plt.plot(grid.X0 + grid.R * np.cos(T), grid.R * np.sin(T), color="black")
# Read grid info
with Dataset(grid_file) as f0:
H = f0.variables["h"][j0:j1, i0:i1]
M = f0.variables["mask_rho"][j0:j1, i0:i1]
lon = f0.variables["lon_rho"][j0:j1, i0:i1]
lat = f0.variables["lat_rho"][j0:j1, i0:i1]
# Cell centers and boundaries
Xcell = np.arange(i0, i1)
Ycell = np.arange(j0, j1)
Xb = np.arange(i0 - 0.5, i1)
Yb = np.arange(j0 - 0.5, j1)
# particle_file
pf = ParticleFile(particle_file)
# Implement the condition, here pid < 500
pids = pf["pid"][t]
# I = pids < 500 # The western half of the particles
# I = pids % 5 == 0 # Alternative condition, every 5th particle
I = np.in1d(pids, [10, 100, 300]) # Explisit subset of pids
# Get the particle positions
X, Y = pf.position(t)
X, Y = X[I], Y[I]
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(1, 1, 1)
# Make background map
def test_array(particle_file):
with ParticleFile(particle_file) as pf:
X = pf.X
assert all(np.array(X) == X.da.values)
assert all(np.array(X[1]) == X[1])
assert all(np.array(X[0:2]) == X[0:2])
# ----------------
f = Dataset(grid_file)
H = f.variables["h"][:, :]
M = f.variables["mask_rho"][:, :]
lon = f.variables["lon_rho"][:, :]
lat = f.variables["lat_rho"][:, :]
f.close()
jmax, imax = H.shape
# ---------------------
# Read particle file
# ---------------------
pf = ParticleFile(particle_file)
# Find record numbers
n0 = -99
# n1 = -99
for n in range(pf.num_times):
if pf.time(n) < date0:
continue
if n0 < 0: # First time
n0 = n
if pf.time(n) < date1:
n1 = n
print("start: ", n0, pf.time(n0))
print("stop : ", n1, pf.time(n1))
def test_variables(particle_file):
"""Indentifies the variables to correct category"""
with ParticleFile(particle_file) as pf:
assert pf.instance_variables == ["pid", "X", "Y"]
assert pf.particle_variables == ["start_time", "location_id"]
with Dataset(grid_file) as f0:
M = f0.variables["mask_rho"][j0:j1, i0:i1]
lon = f0.variables["lon_rho"][j0:j1, i0:i1]
lat = f0.variables["lat_rho"][j0:j1, i0:i1]
# Cell centers and boundaries
Xcell = np.arange(i0, i1)
Ycell = np.arange(j0, j1)
Xb = np.arange(i0 - 0.5, i1)
Yb = np.arange(j0 - 0.5, j1)
# ---------------------------
# Read and count particles
# ---------------------------
pf = ParticleFile(pfile)
X = pf.X[tframe0:tframe1]
Y = pf.Y[tframe0:tframe1]
C = cellcount(X, Y, grid_limits=(i0, i1, j0, j1))
# ------------------------- ---
# Plot particle concentration
# -----------------------------
plt.set_cmap("cool")
plt.pcolormesh(Xb, Yb, C)
plt.colorbar()
# Land mask
constmap = plt.matplotlib.colors.ListedColormap([0.2, 0.6, 0.4])
def test_open():
with pytest.raises(FileNotFoundError):
pf = ParticleFile("no_such_file.nc")