Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Code cleanups #1614

Merged
merged 9 commits into from
Jul 29, 2024
14 changes: 7 additions & 7 deletions parcels/application_kernels/advection.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,21 +251,21 @@ def compute_ds(F0, F1, r, direction, tol):
s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(ds_t / (dxdy * dz)))

# calculate end position in time s_min
def compute_rs(ds, r, B, delta, s_min):
def compute_rs(r, B, delta, s_min):
if abs(B) < tol:
return -delta * s_min + r
else:
return (r + delta / B) * math.exp(-B * s_min) - delta / B

rs_x = compute_rs(ds_x, xsi, B_x, delta_x, s_min)
rs_y = compute_rs(ds_y, eta, B_y, delta_y, s_min)
rs_x = compute_rs(xsi, B_x, delta_x, s_min)
rs_y = compute_rs(eta, B_y, delta_y, s_min)

particle_dlon = (1.-rs_x)*(1.-rs_y) * px[0] + rs_x * (1.-rs_y) * px[1] + rs_x * rs_y * px[2] + (1.-rs_x)*rs_y * px[3] - particle.lon # noqa
particle_dlat = (1.-rs_x)*(1.-rs_y) * py[0] + rs_x * (1.-rs_y) * py[1] + rs_x * rs_y * py[2] + (1.-rs_x)*rs_y * py[3] - particle.lat # noqa
particle_dlon += (1.-rs_x)*(1.-rs_y) * px[0] + rs_x * (1.-rs_y) * px[1] + rs_x * rs_y * px[2] + (1.-rs_x)*rs_y * px[3] - particle.lon # noqa
particle_dlat += (1.-rs_x)*(1.-rs_y) * py[0] + rs_x * (1.-rs_y) * py[1] + rs_x * rs_y * py[2] + (1.-rs_x)*rs_y * py[3] - particle.lat # noqa

if withW:
rs_z = compute_rs(ds_z, zeta, B_z, delta_z, s_min)
particle.depth = (1.-rs_z) * pz[0] + rs_z * pz[1]
rs_z = compute_rs(zeta, B_z, delta_z, s_min)
particle_ddepth += (1.-rs_z) * pz[0] + rs_z * pz[1] - particle.depth # noqa

if particle.dt > 0:
particle.dt = max(direction * s_min * (dxdy * dz), 1e-7)
Expand Down
3 changes: 2 additions & 1 deletion parcels/compilation/codegenerator.py
Original file line number Diff line number Diff line change
Expand Up @@ -657,7 +657,8 @@ def visit_Subscript(self, node):
if isinstance(node.value, FieldNode) or isinstance(node.value, VectorFieldNode):
node.ccode = node.value.__getitem__(node.slice.ccode).ccode
elif isinstance(node.value, ParticleXiYiZiTiAttributeNode):
node.ccode = f"{node.value.obj}->{node.value.attr}[pnum, {node.slice.ccode}]"
ngrid = str(self.fieldset.gridset.size if self.fieldset is not None else 1)
node.ccode = f"{node.value.obj}->{node.value.attr}[pnum*{ngrid}+{node.slice.ccode}]"
elif isinstance(node.value, IntrinsicNode):
raise NotImplementedError(f"Subscript not implemented for object type {type(node.value).__name__}")
else:
Expand Down
2 changes: 0 additions & 2 deletions parcels/kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,6 @@ def __init__(self, fieldset, ptype, pyfunc=None, funcname=None, funccode=None, p
self._ptype = ptype
self._lib = None
self.delete_cfiles = delete_cfiles
self._cleanup_files = None
self._cleanup_lib = None
self._c_include = c_include

# Derive meta information from pyfunc, if not given
Expand Down
35 changes: 19 additions & 16 deletions tests/test_particlefile.py
100644 → 100755
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,7 @@ def Update_lon(particle, fieldset, time):
def test_write_xiyi(fieldset, mode, tmpdir):
outfilepath = tmpdir.join("pfile_xiyi.zarr")
fieldset.U.data[:] = 1 # set a non-zero zonal velocity
fieldset.add_field(Field(name='P', data=np.zeros((2, 20)), lon=np.linspace(0, 1, 20), lat=[0, 2]))
fieldset.add_field(Field(name='P', data=np.zeros((3, 20)), lon=np.linspace(0, 1, 20), lat=[-2, 0, 2]))
dt = 3600

XiYiParticle = ptype[mode].add_variables([
Expand All @@ -306,24 +306,27 @@ def SampleP(particle, fieldset, time):
if time > 5*3600:
tmp = fieldset.P[particle] # noqa

pset = ParticleSet(fieldset, pclass=XiYiParticle, lon=[0], lat=[0.2], lonlatdepth_dtype=np.float64)
pset = ParticleSet(fieldset, pclass=XiYiParticle, lon=[0, 0.2], lat=[0.2, 1], lonlatdepth_dtype=np.float64)
pfile = pset.ParticleFile(name=outfilepath, outputdt=dt)
pset.execute([Get_XiYi, SampleP, AdvectionRK4], endtime=10*dt, dt=dt, output_file=pfile)
pset.execute([SampleP, Get_XiYi, AdvectionRK4], endtime=10*dt, dt=dt, output_file=pfile)

ds = xr.open_zarr(outfilepath)
pxi0 = ds['pxi0'][:].values[0].astype(np.int32)
pxi1 = ds['pxi1'][:].values[0].astype(np.int32)
lons = ds['lon'][:].values[0]
pyi = ds['pyi'][:].values[0].astype(np.int32)
lats = ds['lat'][:].values[0]

assert (pxi0[0] == 0) and (pxi0[-1] == 11) # check that particle has moved
assert np.all(pxi1[:7] == 0) # check that particle has not been sampled on grid 1 until time 6
assert np.all(pxi1[7:] > 0) # check that particle has not been sampled on grid 1 after time 6
for xi, lon in zip(pxi0[1:], lons[1:]):
assert fieldset.U.grid.lon[xi] <= lon < fieldset.U.grid.lon[xi+1]
for yi, lat in zip(pyi[1:], lats[1:]):
assert fieldset.U.grid.lat[yi] <= lat < fieldset.U.grid.lat[yi+1]
pxi0 = ds['pxi0'][:].values.astype(np.int32)
pxi1 = ds['pxi1'][:].values.astype(np.int32)
lons = ds['lon'][:].values
pyi = ds['pyi'][:].values.astype(np.int32)
lats = ds['lat'][:].values

for p in range(pyi.shape[0]):
assert (pxi0[p, 0] == 0) and (pxi0[p, -1] == pset[p].pxi0) # check that particle has moved
assert np.all(pxi1[p, :6] == 0) # check that particle has not been sampled on grid 1 until time 6
assert np.all(pxi1[p, 6:] > 0) # check that particle has not been sampled on grid 1 after time 6
for xi, lon in zip(pxi0[p, 1:], lons[p, 1:]):
assert fieldset.U.grid.lon[xi] <= lon < fieldset.U.grid.lon[xi+1]
for xi, lon in zip(pxi1[p, 6:], lons[p, 6:]):
assert fieldset.P.grid.lon[xi] <= lon < fieldset.P.grid.lon[xi+1]
for yi, lat in zip(pyi[p, 1:], lats[p, 1:]):
assert fieldset.U.grid.lat[yi] <= lat < fieldset.U.grid.lat[yi+1]
ds.close()


Expand Down
Loading