diff --git a/docs/sources/apiref/phyvars.rst b/docs/sources/apiref/phyvars.rst index 9ef59391..8de1e6d3 100644 --- a/docs/sources/apiref/phyvars.rst +++ b/docs/sources/apiref/phyvars.rst @@ -85,19 +85,6 @@ phyvars Dictionary of time series that StagPy can compute. Keys are the variable names, values are :class:`Vart` instances. - .. class:: Varp - - :class:`collections.namedtuple` whose instances hold metadata of - plate variables. It defines the following fields: - - - **description** (*str*): short description of the variable. - - .. data:: PLATES - :annotation: = {platevar: Varp()} - - Dictionary of plate variables output by StagYY. Keys are the variable - names, values are :class:`Varp` instances. - .. data:: SCALES :annotation: = {dimstr: func} diff --git a/docs/sources/apiref/plates.rst b/docs/sources/apiref/plates.rst index c14b0ab1..818f5b8a 100644 --- a/docs/sources/apiref/plates.rst +++ b/docs/sources/apiref/plates.rst @@ -2,4 +2,4 @@ plates ====== .. automodule:: stagpy.plates - :members: cmd + :members: diff --git a/stagpy/commands.py b/stagpy/commands.py index f05ecb2e..0f6feb6f 100644 --- a/stagpy/commands.py +++ b/stagpy/commands.py @@ -142,9 +142,6 @@ def var_cmd(): print('refstate:') _layout(phyvars.REFSTATE, {}) print() - if print_all or conf.var.plates: - print('plates:') - _layout(phyvars.PLATES, {}) def version_cmd(): diff --git a/stagpy/config.py b/stagpy/config.py index acf0a9cf..bd980317 100644 --- a/stagpy/config.py +++ b/stagpy/config.py @@ -179,38 +179,27 @@ def _index_collection(arg): CONF_DEF['plates'] = OrderedDict(( ('plot', - Conf('c,eta,sc', True, 'o', - {'nargs': '?', 'const': '', 'type': str}, - True, 'variables to plot (see stagpy var)')), - ('vzcheck', switch_opt(False, None, - 'activate Colin\'s version with vz checking')), - ('timeprofile', switch_opt(False, None, - 'nb of plates as function of time')), + Conf('c.T.v2-v2.dv2-v2.topo_top', True, 'o', + {'nargs': '?', 'const': '', 'type': str}, True, + 'variables to plot, can be a surface field, field, or dv2')), + ('field', + Conf('eta', True, None, {}, + True, 'field variable to plot with plates info')), + ('stress', switch_opt( + False, None, + 'Plot deviatoric stress instead of velocity on field plots')), + ('continents', switch_opt(True, None, + 'Whether to shade continents on plots')), + ('vzratio', + Conf(0., True, None, {}, True, + 'Ratio of mean vzabs used as threshold for plates limits')), + ('nbplates', switch_opt(False, None, + 'Plot number of plates as function of time')), + ('distribution', switch_opt(False, None, + 'Plot plate size distribution')), ('zoom', Conf(None, True, None, {'type': float}, False, 'zoom around surface')), - ('topomin', Conf(-40, False, None, {}, - True, 'min topography in plots')), - ('topomax', Conf(100, False, None, {}, - True, 'max topography in plots')), - ('agemin', Conf(-50, False, None, {}, - True, 'min age in plots')), - ('agemax', Conf(500, False, None, {}, - True, 'max age in plots')), - ('vmin', Conf(-5000, False, None, {}, - True, 'min velocity in plots')), - ('vmax', Conf(5000, False, None, {}, - True, 'max velocity in plots')), - ('dvmin', Conf(-250000, False, None, {}, - True, 'min velocity derivative in plots')), - ('dvmax', Conf(150000, False, None, {}, - True, 'max velocity derivative in plots')), - ('stressmin', Conf(0, False, None, {}, - True, 'min stress in plots')), - ('stressmax', Conf(800, False, None, {}, - True, 'max stress in plots')), - ('lstressmax', Conf(50, False, None, {}, - True, 'max lithospheric stress in plots')), )) CONF_DEF['info'] = OrderedDict(( @@ -229,8 +218,6 @@ def _index_collection(arg): False, 'print time variables')), ('refstate', Conf(None, True, None, {'action': 'store_true'}, False, 'print refstate variables')), - ('plates', Conf(None, True, None, {'action': 'store_true'}, - False, 'print plates variables')), )) CONF_DEF['config'] = config_conf_section() diff --git a/stagpy/error.py b/stagpy/error.py index 2234df22..1251e26d 100644 --- a/stagpy/error.py +++ b/stagpy/error.py @@ -163,23 +163,6 @@ def __init__(self, zoom): super().__init__(f'Zoom angle should be in [0,360] (received {zoom})') -class StagnantLidError(StagpyError): - """Raised when unexpected stagnant lid regime is found. - - Args: - sdat (:class:`~stagpy.stagyydata.StagyyData`): the StagyyData - instance for which a stagnant lid regime was found. - - Attributes: - sdat (:class:`~stagpy.stagyydata.StagyyData`): the StagyyData - instance for which a stagnant lid regime was found. - """ - - def __init__(self, sdat): - self.sdat = sdat - super().__init__(f'Stagnant lid regime for {sdat}') - - class MissingDataError(StagpyError, KeyError): """Raised when requested data is not present in output.""" diff --git a/stagpy/phyvars.py b/stagpy/phyvars.py index aad37f58..b0b0bd65 100644 --- a/stagpy/phyvars.py +++ b/stagpy/phyvars.py @@ -262,17 +262,6 @@ ('grav', Varr('Gravity', 'Gravity', 'm/s2')), )) -Varp = namedtuple('Varp', ['description']) -PLATES = OrderedDict(( - ('c', Varp('Composition')), - ('eta', Varp('Viscosity')), - ('sc', Varp('Topography')), - ('age', Varp('Age')), - ('str', Varp('Stress')), - ('sx', Varp('Principal deviatoric stress')), - ('ed', Varp('Strain rate')), -)) - SCALES = { 'm': attrgetter('length'), 'kg/m3': attrgetter('density'), diff --git a/stagpy/plates.py b/stagpy/plates.py index 6be943e2..b0cc7934 100644 --- a/stagpy/plates.py +++ b/stagpy/plates.py @@ -1,165 +1,96 @@ """Plate analysis.""" -from contextlib import ExitStack +from contextlib import suppress +from functools import lru_cache import matplotlib.pyplot as plt +from matplotlib import colors import numpy as np -from scipy.signal import argrelextrema +from scipy.signal import argrelmin, argrelmax from . import conf, error, field, phyvars -from ._helpers import saveplot +from ._helpers import saveplot, list_of_vars +from ._step import Field from .stagyydata import StagyyData -def detect_plates_vzcheck(step, vz_thres_ratio=0): - """Detect plates and check with vz and plate size.""" - v_z = step.fields['v3'].values[0, :-1, :, 0] - v_x = step.fields['v2'].values[0, :, :, 0] - tcell = step.fields['T'].values[0, :, :, 0] - n_z = step.geom.nztot - nphi = step.geom.nptot - r_c = step.geom.r_centers - r_w = step.geom.r_walls - dphi = 1 / nphi - - flux_c = n_z * [0] - for i_z in range(0, n_z): - flux_c[i_z] = np.sum((tcell[:, i_z] - step.timeinfo.loc['Tmean']) * - v_z[:, i_z]) * r_w[i_z] * dphi - - # checking stagnant lid, criterion seems weird! - if all(abs(flux_c[i_z]) <= np.max(flux_c) / 50 - for i_z in range(n_z - n_z // 20, n_z)): - raise error.StagnantLidError(step.sdat) - - # verifying horizontal plate speed and closeness of plates - vphi_surf = v_x[:, -1] - dvphi = np.diff(vphi_surf) / (r_c[-1] * dphi) - dvx_thres = 16 * step.timeinfo.loc['vrms'] - - limits = [ - phi for phi in range(nphi) - if (abs(dvphi[phi]) >= dvx_thres and - all(abs(dvphi[i % nphi]) <= abs(dvphi[phi]) - for i in range(phi - nphi // 33, phi + nphi // 33))) - ] - +def _vzcheck(iphis, snap, vz_thres): + """Remove positions where vz is below threshold.""" # verifying vertical velocity - vz_thres = 0 - if vz_thres_ratio > 0: - vz_mean = (np.sum(step.rprofs['vzabs'].values * np.diff(r_w)) / - (r_w[-1] - r_w[0])) - vz_thres = vz_mean * vz_thres_ratio - k = 0 - for i in range(len(limits)): - vzm = 0 - phi = limits[i - k] - for i_z in range(1 if phi == nphi - 1 else 0, n_z): - vzm += (abs(v_z[phi, i_z]) + - abs(v_z[phi - 1, i_z]) + - abs(v_z[(phi + 1) % nphi, i_z])) / (n_z * 3) - + vzabs = np.abs(snap.fields['v3'].values[0, ..., 0]) + argdel = [] + for i, iphi in enumerate(iphis): + vzm = np.mean(vzabs[[iphi - 1, iphi, iphi + 1], :]) if vzm < vz_thres: - limits.remove(phi) - k += 1 + argdel.append(i) + return np.delete(iphis, argdel) - return limits, dvphi, vphi_surf +@lru_cache() +def detect_plates(snap, vz_thres_ratio=0): + """Detect plate limits using derivative of horizontal velocity. -def detect_plates(step, vrms_surface, fids, time): - """Detect plates using derivative of horizontal velocity.""" - vphi = step.fields['v2'].values[0, :, :, 0] - ph_coord = step.geom.p_centers + This function is cached for convenience. - if step.sdat.par['boundaries']['air_layer']: - dsa = step.sdat.par['boundaries']['air_thickness'] - # we are a bit below the surface; should check if you are in the - # thermal boundary layer - indsurf = np.argmin( - np.abs(1 - dsa - step.geom.r_centers + step.geom.rcmb)) - 4 - else: - indsurf = -1 - - # vphi at cell-center - vph2 = 0.5 * (vphi[1:] + vphi[:-1]) - # dvphi/dphi at cell-center - dvph2 = np.diff(vphi[:, indsurf]) / (ph_coord[1] - ph_coord[0]) - - io_surface(step.isnap, time, fids[6], dvph2) - io_surface(step.isnap, time, fids[7], vph2[:, indsurf]) + Args: + snap (:class:`~stagpy._step.Step`): a step of a StagyyData instance. + vz_thres_ratio (float): if above zero, an addition check based on the + vertical velocities is performed. Limits detected above a region + where the vertical velocity is below vz_thres_ratio * mean(vzabs) + are ignored. + Returns: + tuple of :class:`numpy.array`: itrenches, iridges + 1D arrays containing phi-index of detected trenches and ridges. + """ + dvphi = _surf_diag(snap, 'dv2').values # finding trenches - pom2 = np.copy(dvph2) - maskbigdvel = -vrms_surface * ( - 30 if step.sdat.par['boundaries']['air_layer'] else 10) - pom2[pom2 > maskbigdvel] = maskbigdvel - trench_span = 15 if step.sdat.par['boundaries']['air_layer'] else 10 - argless_dv = argrelextrema( - pom2, np.less, order=trench_span, mode='wrap')[0] - trench = ph_coord[argless_dv] - velocity_trench = vph2[argless_dv, indsurf] - dv_trench = dvph2[argless_dv] + dvphi_saturated = np.copy(dvphi) + max_dvphi = np.amin(dvphi) * 0.2 + dvphi_saturated[dvphi > max_dvphi] = max_dvphi + trench_span = 15 if snap.sdat.par['boundaries']['air_layer'] else 10 + itrenches = argrelmin(dvphi_saturated, order=trench_span, mode='wrap')[0] # finding ridges - pom2 = np.copy(dvph2) - masksmalldvel = np.amax(dvph2) * 0.2 - pom2[pom2 < masksmalldvel] = masksmalldvel + dvphi_saturated = np.copy(dvphi) + min_dvphi = np.amax(dvphi) * 0.2 + dvphi_saturated[dvphi < min_dvphi] = min_dvphi ridge_span = 20 - arggreat_dv = argrelextrema( - pom2, np.greater, order=ridge_span, mode='wrap')[0] - ridge = ph_coord[arggreat_dv] + iridges = argrelmax(dvphi_saturated, order=ridge_span, mode='wrap')[0] - # elimination of ridges that are too close to trench + # elimination of ridges that are too close to a trench + phi = snap.geom.p_centers + phi_trenches = phi[itrenches] argdel = [] - if trench.any() and ridge.any(): - for i, ridge_i in enumerate(ridge): - mdistance = np.amin(abs(trench - ridge_i)) + if itrenches.size and iridges.size: + for i, iridge in enumerate(iridges): + mdistance = np.amin(np.abs(phi_trenches - phi[iridge])) if mdistance < 0.016: argdel.append(i) if argdel: - print('deleting from ridge', trench, ridge[argdel]) - ridge = np.delete(ridge, np.array(argdel)) - arggreat_dv = np.delete(arggreat_dv, np.array(argdel)) - - dv_ridge = dvph2[arggreat_dv] - if 'age' in conf.plates.plot: - agefld = step.fields['age'].values[0, :, :, 0] - age_surface = np.ma.masked_where(agefld[:, indsurf] < 0.00001, - agefld[:, indsurf]) - age_surface_dim = age_surface * vrms_surface *\ - conf.scaling.ttransit / conf.scaling.yearins / 1.e6 - agetrench = age_surface_dim[argless_dv] # age at the trench - else: - agetrench = np.zeros(len(argless_dv)) + iridges = np.delete(iridges, argdel) - # writing the output into a file, all time steps are in one file - for itrench in np.arange(len(trench)): - fids[0].write("%7.0f %11.7f %10.6f %9.2f %9.2f \n" % ( - step.isnap, - step.time, - trench[itrench], - velocity_trench[itrench], - agetrench[itrench] - )) + # additional check on vz + if vz_thres_ratio > 0: + r_w = snap.geom.r_walls + vz_mean = (np.sum(snap.rprofs['vzabs'].values * np.diff(r_w)) / + (r_w[-1] - r_w[0])) + vz_thres = vz_mean * vz_thres_ratio + itrenches = _vzcheck(itrenches, snap, vz_thres) + iridges = _vzcheck(iridges, snap, vz_thres) - return trench, ridge, agetrench, dv_trench, dv_ridge + return itrenches, iridges -def plot_plate_limits(axis, ridges, trenches, ymin, ymax): +def _plot_plate_limits(axis, trenches, ridges): """Plot lines designating ridges and trenches.""" for trench in trenches: - axis.axvline( - x=trench, ymin=ymin, ymax=ymax, - color='red', ls='dashed', alpha=0.4) + axis.axvline(x=trench, color='red', ls='dashed', alpha=0.4) for ridge in ridges: - axis.axvline( - x=ridge, ymin=ymin, ymax=ymax, - color='green', ls='dashed', alpha=0.4) - axis.set_xlim(0, 2 * np.pi) - axis.set_ylim(ymin, ymax) + axis.axvline(x=ridge, color='green', ls='dashed', alpha=0.4) -def plot_plate_limits_field(axis, rcmb, ridges, trenches): +def _plot_plate_limits_field(axis, rcmb, trenches, ridges): """Plot arrows designating ridges and trenches in 2D field plots.""" for trench in trenches: xxd = (rcmb + 1.02) * np.cos(trench) # arrow begin @@ -167,566 +98,224 @@ def plot_plate_limits_field(axis, rcmb, ridges, trenches): xxt = (rcmb + 1.35) * np.cos(trench) # arrow end yyt = (rcmb + 1.35) * np.sin(trench) # arrow end axis.annotate('', xy=(xxd, yyd), xytext=(xxt, yyt), - arrowprops=dict(facecolor='red', shrink=0.05)) + arrowprops=dict(facecolor='red', shrink=0.05), + annotation_clip=False) for ridge in ridges: xxd = (rcmb + 1.02) * np.cos(ridge) yyd = (rcmb + 1.02) * np.sin(ridge) xxt = (rcmb + 1.35) * np.cos(ridge) yyt = (rcmb + 1.35) * np.sin(ridge) axis.annotate('', xy=(xxd, yyd), xytext=(xxt, yyt), - arrowprops=dict(facecolor='green', shrink=0.05)) - - -def plot_plates(step, time, vrms_surface, trench, ridge, agetrench, - topo, fids): - """Handle plotting stuff.""" - vphi = step.fields['v2'].values[0, :, :, 0] - tempfld = step.fields['T'].values[0, :, :, 0] - concfld = step.fields['c'].values[0, :, :, 0] - timestep = step.isnap - - if step.sdat.par['boundaries']['air_layer']: - dsa = step.sdat.par['boundaries']['air_thickness'] - # we are a bit below the surface; delete "-some number" - # to be just below - # the surface (that is considered plane here); should check if you are + arrowprops=dict(facecolor='green', shrink=0.05), + annotation_clip=False) + + +def _isurf(snap): + """Return index of surface accounting for air layer.""" + if snap.sdat.par['boundaries']['air_layer']: + dsa = snap.sdat.par['boundaries']['air_thickness'] + # we are a bit below the surface; delete "-some number" to be just + # below the surface (that is considered plane here); should check if # in the thermal boundary layer - indsurf = np.argmin( - np.abs(1 - dsa - step.geom.r_centers + step.geom.rcmb)) - 4 - # depth to detect the continents - indcont = np.argmin( - np.abs(1 - dsa - step.geom.r_centers + step.geom.rcmb)) - 10 + isurf = np.argmin( + np.abs(1 - dsa - snap.geom.r_centers + snap.geom.rcmb)) - 4 else: - indsurf = -1 - indcont = -1 # depth to detect continents - - if step.sdat.par['boundaries']['air_layer'] and\ - not step.sdat.par['continents']['proterozoic_belts']: - continents = np.ma.masked_where( - np.logical_or(concfld[:, indcont] < 3, - concfld[:, indcont] > 4), - concfld[:, indcont]) - elif (step.sdat.par['boundaries']['air_layer'] and - step.sdat.par['continents']['proterozoic_belts']): - continents = np.ma.masked_where( - np.logical_or(concfld[:, indcont] < 3, - concfld[:, indcont] > 5), - concfld[:, indcont]) - elif step.sdat.par['tracersin']['tracers_weakcrust']: - continents = np.ma.masked_where( - concfld[:, indcont] < 3, concfld[:, indcont]) + isurf = -1 + return isurf + + +def _surf_diag(snap, name): + """Get a surface field. + + Can be a sfield, a regular scalar field evaluated at the surface, + or dv2 (which is dvphi/dphi). + """ + with suppress(error.UnknownVarError): + return snap.sfields[name] + isurf = _isurf(snap) + with suppress(error.UnknownVarError): + field, meta = snap.fields[name] + return Field(field[0, :, isurf, 0], meta) + if name == 'dv2': + vphi = snap.fields['v2'].values[0, :, isurf, 0] + dvphi = np.diff(vphi) / (snap.geom.r_centers[isurf] * + np.diff(snap.geom.p_walls)) + return Field(dvphi, phyvars.Varf(r"$dv_\phi/d\phi$", '1/s')) + raise error.UnknownVarError(name) + + +def _continents_location(snap, at_surface=True): + """Location of continents as a boolean array. + + If at_surface is True, it is evaluated only at the surface, otherwise it is + evaluated in the entire domain. + """ + if at_surface: + if snap.sdat.par['boundaries']['air_layer']: + icont = _isurf(snap) - 6 + else: + icont = -1 + else: + icont = slice(None) + csurf = snap.fields['c'].values[0, :, icont, 0] + if snap.sdat.par['boundaries']['air_layer'] and\ + not snap.sdat.par['continents']['proterozoic_belts']: + return (csurf >= 3) & (csurf <= 4) + elif (snap.sdat.par['boundaries']['air_layer'] and + snap.sdat.par['continents']['proterozoic_belts']): + return (csurf >= 3) & (csurf <= 5) + elif snap.sdat.par['tracersin']['tracers_weakcrust']: + return csurf >= 3 + return csurf >= 2 + + +def plot_at_surface(snap, names): + """Plot surface diagnostics. + + Args: + snap (:class:`~stagpy._step.Step`): a step of a StagyyData instance. + names (str): names of requested surface diagnotics. They are separated + by ``-`` (figures), ``.`` (subplots) and ``,`` (same subplot). + Surface diagnotics can be valid surface field names, field names, + or `"dv2"` which is d(vphi)/dphi. + """ + for vfig in list_of_vars(names): + fig, axes = plt.subplots(nrows=len(vfig), sharex=True, + figsize=(12, 2 * len(vfig))) + axes = [axes] if len(vfig) == 1 else axes + fname = 'plates_surf_' + for axis, vplt in zip(axes, vfig): + fname += '_'.join(vplt) + '_' + label = '' + for name in vplt: + data, meta = _surf_diag(snap, name) + label = meta.description + phi = (snap.geom.p_centers if data.size == snap.geom.nptot + else snap.geom.p_walls) + axis.plot(phi, data, label=label) + axis.set_ylim([conf.plot.vmin, conf.plot.vmax]) + if conf.plates.continents: + continents = _continents_location(snap) + ymin, ymax = axis.get_ylim() + axis.fill_between(snap.geom.p_centers, ymin, ymax, + where=continents, alpha=0.2, + facecolor='#8B6914') + axis.set_ylim([ymin, ymax]) + phi = snap.geom.p_centers + itrenches, iridges = detect_plates(snap, conf.plates.vzratio) + _plot_plate_limits(axis, phi[itrenches], phi[iridges]) + if len(vplt) == 1: + axis.set_ylabel(label) + else: + axis.legend() + axes[-1].set_xlabel(r"$\phi$") + axes[-1].set_xlim(snap.geom.p_walls[[0, -1]]) + saveplot(fig, fname, snap.isnap) + + +def _write_trench_diagnostics(step, vrms_surf, fid): + """Print out some trench diagnostics.""" + itrenches, _ = detect_plates(step, conf.plates.vzratio) + time = step.time * vrms_surf *\ + conf.scaling.ttransit / conf.scaling.yearins / 1.e6 + isurf = _isurf(step) + trenches = step.geom.p_centers[itrenches] + + # vphi at trenches + vphi = step.fields['v2'].values[0, :, isurf, 0] + vphi = (vphi[1:] + vphi[:-1]) / 2 + v_trenches = vphi[itrenches] + + if 'age' in step.fields: + agefld = step.fields['age'].values[0, :, isurf, 0] + age_surface = np.ma.masked_where(agefld < 1.e-5, agefld) + age_surface_dim = age_surface * vrms_surf *\ + conf.scaling.ttransit / conf.scaling.yearins / 1.e6 + agetrenches = age_surface_dim[itrenches] # age at the trench else: - continents = np.ma.masked_where( - concfld[:, indcont] < 2, concfld[:, indcont]) - - # masked array, only continents are true - continentsall = continents / continents - - ph_coord = step.geom.p_centers - - # velocity derivative at cell-center - dvph2 = np.diff(vphi[:, indsurf]) / (ph_coord[1] - ph_coord[0]) - - # plotting - fig0, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True, figsize=(12, 8)) - ax1.plot(ph_coord, concfld[:, indsurf], color='g', label='Conc') - ax2.plot(ph_coord, tempfld[:, indsurf], color='k', label='Temp') - ax3.plot(step.geom.p_walls, vphi[:, indsurf], label='Vel') - - ax1.fill_between( - ph_coord, continents, 1., facecolor='#8B6914', alpha=0.2) - ax2.fill_between( - ph_coord, continentsall, 0., facecolor='#8B6914', alpha=0.2) - - tempmin = step.sdat.par['boundaries']['topT_val'] * 0.9\ - if step.sdat.par['boundaries']['topT_mode'] == 'iso' else 0.0 - tempmax = step.sdat.par['boundaries']['botT_val'] * 0.35\ - if step.sdat.par['boundaries']['botT_mode'] == 'iso' else 0.8 - - ax2.set_ylim(tempmin, tempmax) - ax3.fill_between( - ph_coord, continentsall * round(1.5 * np.amax(dvph2), 1), - round(np.amin(dvph2) * 1.1, 1), facecolor='#8B6914', alpha=0.2) - ax3.set_ylim(conf.plates.vmin, conf.plates.vmax) - - ax1.set_ylabel("Concentration") - ax2.set_ylabel("Temperature") - ax3.set_ylabel("Velocity") - ax1.set_title(timestep) - ax1.text(0.95, 1.07, str(round(time, 0)) + ' My', - transform=ax1.transAxes) - ax1.text(0.01, 1.07, str(round(step.time, 8)), - transform=ax1.transAxes) - - plot_plate_limits(ax3, ridge, trench, conf.plates.vmin, - conf.plates.vmax) - - saveplot(fig0, 'sveltempconc', timestep) - - # plotting velocity and velocity derivative - fig0, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(12, 8)) - ax1.plot(step.geom.p_walls, vphi[:, indsurf], label='Vel') - ax1.axhline(y=0, xmin=0, xmax=2 * np.pi, - color='black', ls='solid', alpha=0.2) - ax1.set_ylabel("Velocity") - ax1.text(0.95, 1.07, str(round(time, 0)) + ' My', - transform=ax1.transAxes) - ax1.text(0.01, 1.07, str(round(step.time, 8)), - transform=ax1.transAxes) - ax2.plot(ph_coord, dvph2, color='k', label='dv') - ax2.set_ylabel("dv") - - plot_plate_limits(ax1, ridge, trench, conf.plates.vmin, - conf.plates.vmax) - plot_plate_limits(ax2, ridge, trench, conf.plates.dvmin, - conf.plates.dvmax) - ax1.set_xlim(0, 2 * np.pi) - ax1.set_title(timestep) - - ax1.fill_between( - ph_coord, continentsall * conf.plates.vmin, conf.plates.vmax, - facecolor='#8b6914', alpha=0.2) - ax1.set_ylim(conf.plates.vmin, conf.plates.vmax) - ax2.fill_between( - ph_coord, continentsall * conf.plates.dvmin, - conf.plates.dvmax, facecolor='#8b6914', alpha=0.2) - ax2.set_ylim(conf.plates.dvmin, conf.plates.dvmax) - - saveplot(fig0, 'sveldvel', timestep) - - # plotting velocity and second invariant of stress - if 'str' in conf.plates.plot: - stressfld = step.fields['sII'].values[0, :, :, 0] - fig0, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(12, 8)) - ax1.plot(step.geom.p_walls, vphi[:, indsurf], label='Vel') - ax1.axhline(y=0, xmin=0, xmax=2 * np.pi, - color='black', ls='solid', alpha=0.2) - ax1.set_ylabel("Velocity") - ax1.text(0.95, 1.07, str(round(time, 0)) + ' My', - transform=ax1.transAxes) - ax1.text(0.01, 1.07, str(round(step.time, 8)), - transform=ax1.transAxes) - ax2.plot(ph_coord, - stressfld[:, indsurf] * step.sdat.scales.stress / 1.e6, - color='k', label='Stress') - ax2.set_ylim(conf.plates.stressmin, conf.plates.stressmax) - ax2.set_ylabel("Stress [MPa]") - - plot_plate_limits(ax1, ridge, trench, - conf.plates.vmin, conf.plates.vmax) - plot_plate_limits(ax2, ridge, trench, - conf.plates.stressmin, conf.plates.stressmax) - ax1.set_xlim(0, 2 * np.pi) - ax1.set_title(timestep) - - ax1.fill_between( - ph_coord, continentsall * conf.plates.vmin, - conf.plates.vmax, facecolor='#8B6914', alpha=0.2) - ax1.set_ylim(conf.plates.vmin, conf.plates.vmax) - ax2.fill_between( - ph_coord, continentsall * conf.plates.dvmin, - conf.plates.dvmax, - facecolor='#8B6914', alpha=0.2) - - saveplot(fig0, 'svelstress', timestep) - - # plotting velocity - fig1, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(12, 8)) - ax1.plot(step.geom.p_walls, vphi[:, indsurf], label='Vel') - ax1.axhline(y=0, xmin=0, xmax=2 * np.pi, - color='black', ls='solid', alpha=0.2) - ax1.set_ylim(conf.plates.vmin, conf.plates.vmax) - ax1.set_ylabel("Velocity") - ax1.text(0.95, 1.07, str(round(time, 0)) + ' My', - transform=ax1.transAxes) - plot_plate_limits(ax1, ridge, trench, conf.plates.vmin, - conf.plates.vmax) - - # plotting velocity and age at surface - if 'age' in conf.plates.plot: - agefld = step.fields['age'].values[0, :, :, 0] - age_surface = np.ma.masked_where( - agefld[:, indsurf] < 0.00001, agefld[:, indsurf]) - age_surface_dim = (age_surface * vrms_surface * conf.scaling.ttransit / - conf.scaling.yearins / 1.e6) - - fig2, (ax3, ax4) = plt.subplots(2, 1, sharex=True, figsize=(12, 8)) - ax3.plot(ph_coord, vphi[:, indsurf], label='Vel') - ax3.axhline( - y=0, xmin=0, xmax=2 * np.pi, - color='black', ls='solid', alpha=0.2) - ax3.set_ylim(conf.plates.vmin, conf.plates.vmax) - ax3.set_ylabel("Velocity") - ax3.text(0.95, 1.07, str(round(time, 0)) + ' My', - transform=ax3.transAxes) - ax3.fill_between( - ph_coord, continentsall * conf.plates.vmax, - conf.plates.vmin, facecolor='#8B6914', alpha=0.2) - plot_plate_limits(ax3, ridge, trench, - conf.plates.vmin, conf.plates.vmax) - - times_subd = [] - age_subd = [] + agetrenches = np.zeros(len(itrenches)) + + if conf.plates.continents: + phi_cont = step.geom.p_centers[_continents_location(step)] + else: + phi_cont = np.array([np.nan]) + distance_subd = [] - ph_trench_subd = [] ph_cont_subd = [] - if step.sdat.par['switches']['cont_tracers']: - for i, trench_i in enumerate(trench): - # detection of the distance in between subduction and continent - angdistance1 = abs(ph_coord[continentsall == 1] - trench_i) - angdistance2 = 2. * np.pi - angdistance1 - angdistance = np.minimum(angdistance1, angdistance2) - distancecont = min(angdistance) - argdistancecont = np.argmin(angdistance) - continentpos = ph_coord[continentsall == 1][argdistancecont] - - ph_trench_subd.append(trench_i) - age_subd.append(agetrench[i]) - ph_cont_subd.append(continentpos) - distance_subd.append(distancecont) - times_subd.append(step.time) - - if angdistance1[argdistancecont] < angdistance2[argdistancecont]: - if continentpos - trench_i < 0: # continent is on the left - distancecont = - distancecont - ax1.annotate('', xy=(trench_i + distancecont, 2000), - xycoords='data', xytext=(trench_i, 2000), - textcoords='data', - arrowprops=dict(arrowstyle="->", lw="2", - shrinkA=0, shrinkB=0)) - else: # distance over boundary - xy_anot, xy_text = 0, 2 * np.pi - if continentpos - trench_i < 0: - xy_anot, xy_text = xy_text, xy_anot - ax1.annotate('', xy=(xy_anot, 2000), - xycoords='data', xytext=(trench_i, 2000), - textcoords='data', - arrowprops=dict(arrowstyle="-", lw="2", - shrinkA=0, shrinkB=0)) - ax1.annotate('', xy=(continentpos, 2000), - xycoords='data', xytext=(xy_text, 2000), - textcoords='data', - arrowprops=dict(arrowstyle="->", lw="2", - shrinkA=0, shrinkB=0)) - - ax1.fill_between( - ph_coord, continentsall * conf.plates.vmin, - conf.plates.vmax, facecolor='#8B6914', alpha=0.2) - ax2.set_ylabel("Topography [km]") - ax2.axhline(y=0, xmin=0, xmax=2 * np.pi, - color='black', ls='solid', alpha=0.2) - ax2.plot(topo[:, 0], - topo[:, 1] * step.sdat.scales.length / 1.e3, - color='black') - ax2.set_xlim(0, 2 * np.pi) - ax2.set_ylim(conf.plates.topomin, conf.plates.topomax) - ax2.fill_between( - ph_coord, continentsall * conf.plates.topomax, - conf.plates.topomin, facecolor='#8B6914', alpha=0.2) - plot_plate_limits(ax2, ridge, trench, conf.plates.topomin, - conf.plates.topomax) - ax1.set_title(timestep) - saveplot(fig1, 'sveltopo', timestep) - - if 'age' in conf.plates.plot: - ax4.set_ylabel("Seafloor age [My]") - # in dimensions - ax4.plot(ph_coord, age_surface_dim, color='black') - ax4.set_xlim(0, 2 * np.pi) - ax4.fill_between( - ph_coord, continentsall * conf.plates.agemax, - conf.plates.agemin, facecolor='#8B6914', alpha=0.2) - ax4.set_ylim(conf.plates.agemin, conf.plates.agemax) - plot_plate_limits(ax4, ridge, trench, conf.plates.agemin, - conf.plates.agemax) - ax3.set_title(timestep) - saveplot(fig2, 'svelage', timestep) + for trench_i in trenches: + # compute distance between subduction and continent + angdistance1 = np.abs(phi_cont - trench_i) + angdistance2 = 2 * np.pi - angdistance1 + angdistance = np.minimum(angdistance1, angdistance2) + i_closest = np.argmin(angdistance) + + ph_cont_subd.append(phi_cont[i_closest]) + distance_subd.append(angdistance[i_closest]) # writing the output into a file, all time steps are in one file - for isubd in np.arange(len(distance_subd)): - fids[1].write("%6.0f %11.7f %11.3f %10.6f %10.6f %10.6f %11.3f\n" % ( - timestep, - times_subd[isubd], - time, - distance_subd[isubd], - ph_trench_subd[isubd], - ph_cont_subd[isubd], - age_subd[isubd], - )) - - -def io_surface(timestep, time, fid, fld): - """Output surface files.""" - fid.write(f"{timestep} {time}") - fid.writelines(["%10.2e" % item for item in fld[:]]) - fid.writelines(["\n"]) - - -def lithospheric_stress(step, trench, ridge, time): - """Calculate stress in the lithosphere.""" - timestep = step.isnap - base_lith = step.geom.rcmb + 1 - 0.105 - - stressfld = step.fields['sII'].values[0, :, :, 0] - r_centers = np.outer(np.ones(stressfld.shape[0]), step.geom.r_centers) - stressfld = np.ma.masked_where(r_centers < base_lith, stressfld) - - # stress integration in the lithosphere - dzm = (step.geom.r_centers[1:] - step.geom.r_centers[:-1]) - stress_lith = np.sum((stressfld[:, 1:] * dzm.T), axis=1) - ph_coord = step.geom.p_centers # probably doesn't need alias - - # plot stress in the lithosphere - fig, axis, _, _ = field.plot_scalar(step, 'sII', stressfld, - cmap='plasma_r', vmin=0, vmax=300) - # Annotation with time and step - axis.text(1., 0.9, str(round(time, 0)) + ' My', transform=axis.transAxes) - axis.text(1., 0.1, str(timestep), transform=axis.transAxes) - saveplot(fig, 'lith', timestep) - - # velocity - vphi = step.fields['v2'].values[0, :, :, 0] - - # position of continents - concfld = step.fields['c'].values[0, :, :, 0] - if step.sdat.par['boundaries']['air_layer']: - # we are a bit below the surface; delete "-some number" - # to be just below - dsa = step.sdat.par['boundaries']['air_thickness'] - # depth to detect the continents - indcont = np.argmin( - np.abs(1 - dsa - step.geom.r_centers + step.geom.rcmb)) - 10 - else: - # depth to detect continents - indcont = -1 - if step.sdat.par['boundaries']['air_layer'] and\ - not step.sdat.par['continents']['proterozoic_belts']: - continents = np.ma.masked_where( - np.logical_or(concfld[:-1, indcont] < 3, - concfld[:-1, indcont] > 4), - concfld[:-1, indcont]) - elif step.sdat.par['boundaries']['air_layer'] and\ - step.sdat.par['continents']['proterozoic_belts']: - continents = np.ma.masked_where( - np.logical_or(concfld[:-1, indcont] < 3, - concfld[:-1, indcont] > 5), - concfld[:-1, indcont]) - elif step.sdat.par['tracersin']['tracers_weakcrust']: - continents = np.ma.masked_where( - concfld[:-1, indcont] < 3, concfld[:-1, indcont]) - else: - continents = np.ma.masked_where( - concfld[:-1, indcont] < 2, concfld[:-1, indcont]) - - # masked array, only continents are true - continentsall = continents / continents - - # plot integrated stress in the lithosphere - fig0, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(12, 8)) - ax1.plot(step.geom.p_walls, vphi[:, -1], label='Vel') - ax1.axhline(y=0, xmin=0, xmax=2 * np.pi, - color='black', ls='solid', alpha=0.2) - ax1.set_ylabel("Velocity") - ax1.text(0.95, 1.07, str(round(time, 0)) + ' My', - transform=ax1.transAxes) - ax1.text(0.01, 1.07, str(round(step.time, 8)), - transform=ax1.transAxes) - - intstr_scale = step.sdat.scales.stress * step.sdat.scales.length / 1.e12 - ax2.plot(ph_coord, stress_lith * intstr_scale, color='k', label='Stress') - ax2.set_ylabel(r"Integrated stress [$TN\,m^{-1}$]") - - plot_plate_limits(ax1, ridge, trench, conf.plates.vmin, - conf.plates.vmax) - plot_plate_limits(ax2, ridge, trench, conf.plates.stressmin, - conf.plates.lstressmax) - ax1.set_xlim(0, 2 * np.pi) - ax1.set_title(timestep) - - ax1.fill_between( - ph_coord, continentsall * conf.plates.vmin, - conf.plates.vmax, facecolor='#8b6914', alpha=0.2) - ax1.set_ylim(conf.plates.vmin, conf.plates.vmax) - ax2.fill_between( - ph_coord, continentsall * conf.plates.stressmin, - conf.plates.lstressmax, facecolor='#8b6914', alpha=0.2) - ax2.set_ylim(conf.plates.stressmin, conf.plates.lstressmax) - - saveplot(fig0, 'svelslith', timestep) - - -def set_of_vars(arg_plot): - """Build set of needed variables. + for isubd in range(len(trenches)): + fid.write( + "%6.0f %11.7f %11.3f %10.6f %10.6f %10.6f %10.6f %11.3f\n" % ( + step.isnap, + step.time, + time, + trenches[isubd], + v_trenches[isubd], + distance_subd[isubd], + ph_cont_subd[isubd], + agetrenches[isubd])) + + +def plot_scalar_field(snap, fieldname): + """Plot scalar field with plate information. Args: - arg_plot (str): string with variable names separated with ``,``. - Returns: - set of str: set of variables. + snap (:class:`~stagpy._step.Step`): a step of a StagyyData instance. + fieldname (str): name of the field that should be decorated with plate + informations. """ - return set(var for var in arg_plot.split(',') if var in phyvars.PLATES) - - -def main_plates(sdat): - """Plot several plates information.""" - # averaged horizontal surface velocity needed for redimensionalisation - uprof_averaged, radius, _ = sdat.walk.filter(rprofs=True)\ - .rprofs_averaged['vhrms'] - if sdat.par['boundaries']['air_layer']: - dsa = sdat.par['boundaries']['air_thickness'] - isurf = np.argmin(abs(radius - radius[-1] + dsa)) - vrms_surface = uprof_averaged.iloc[isurf] - isurf = np.argmin(abs((1 - dsa) - radius)) - isurf -= 4 # why different isurf for the rest? - else: - isurf = -1 - vrms_surface = uprof_averaged[isurf] - - # determine names of files - fnames = ['plate_velocity', 'distance_subd', 'continents', - 'flux', 'topography', 'age', 'velderiv', 'velocity'] - fnames = [f'plates_{stem}_{sdat.walk.stepstr}' for stem in fnames] - with ExitStack() as stack: - fids = [stack.enter_context(open(fname, 'w')) for fname in fnames] - fids[0].write('# it time ph_trench vel_trench age_trench\n') - fids[1].write('# it time time [My] distance ' - 'ph_trench ph_cont age_trench [My]\n') - - for step in sdat.walk.filter(fields=['T']): - # could check other fields too - timestep = step.isnap - print('Treating snapshot', timestep) - - rcmb = step.geom.rcmb - # topography - fname = sdat.filename('sc', timestep=timestep, suffix='.dat') - topo = np.genfromtxt(str(fname)) - # rescaling topography! - if sdat.par['boundaries']['air_layer']: - topo[:, 1] = topo[:, 1] / (1. - dsa) - - time = step.time * vrms_surface *\ - conf.scaling.ttransit / conf.scaling.yearins / 1.e6 - trenches, ridges, agetrenches, _, _ =\ - detect_plates(step, vrms_surface, fids, time) - plot_plates(step, time, vrms_surface, trenches, ridges, - agetrenches, topo, fids) - - # prepare for continent plotting - concfld = step.fields['c'].values[0, :, :, 0] - continentsfld = np.ma.masked_where( - concfld < 3, concfld) # plotting continents, to-do - continentsfld = continentsfld / continentsfld - - temp = step.fields['T'].values[0, :, :, 0] - tgrad = (temp[:, isurf - 1] - temp[:, isurf]) /\ - (step.geom.r_centers[isurf] - step.geom.r_centers[isurf - 1]) - - io_surface(timestep, time, fids[2], concfld[:-1, isurf]) - io_surface(timestep, time, fids[3], tgrad) - io_surface(timestep, time, fids[4], topo[:, 1]) - if 'age' in conf.plates.plot: - io_surface(timestep, time, fids[5], - step.fields['age'].values[0, :, isurf, 0]) - - # plot viscosity field with position of trenches and ridges - etamin, _ = sdat.scale(1e-2, 'Pa') - etamax, _ = sdat.scale(sdat.par['viscosity']['eta_max'], 'Pa') - fig, axis, _, _ = field.plot_scalar(step, 'eta', - vmin=etamin, vmax=etamax) - - # plotting continents - cbar = conf.field.colorbar - conf.field.colorbar = False - field.plot_scalar(step, 'c', continentsfld, axis, - cmap='cool_r', vmin=0, vmax=0) - cmap2 = plt.cm.ocean - cmap2.set_over('m') - conf.field.colorbar = cbar - - # plotting velocity vectors - field.plot_vec(axis, step, 'v') - - # Annotation with time and step - axis.text(1., 0.9, str(round(time, 0)) + ' My', - transform=axis.transAxes) - axis.text(1., 0.1, str(timestep), - transform=axis.transAxes) - - # Put arrow where ridges and trenches are - plot_plate_limits_field(axis, rcmb, ridges, trenches) - - saveplot(fig, 'eta', timestep, close=conf.plates.zoom is None) - - # Zoom - if conf.plates.zoom is not None: - if not 0 <= conf.plates.zoom <= 360: - raise error.InvalidZoomError(conf.plates.zoom) - if 45 < conf.plates.zoom <= 135: - ladd, radd, uadd, dadd = 0.8, 0.8, 0.05, 0.1 - elif 135 < conf.plates.zoom <= 225: - ladd, radd, uadd, dadd = 0.05, 0.1, 0.8, 0.8 - elif 225 < conf.plates.zoom <= 315: - ladd, radd, uadd, dadd = 0.8, 0.8, 0.1, 0.05 - else: # >315 or <=45 - ladd, radd, uadd, dadd = 0.1, 0.05, 0.8, 0.8 - xzoom = (rcmb + 1) * np.cos(np.radians(conf.plates.zoom)) - yzoom = (rcmb + 1) * np.sin(np.radians(conf.plates.zoom)) - axis.set_xlim(xzoom - ladd, xzoom + radd) - axis.set_ylim(yzoom - dadd, yzoom + uadd) - saveplot(fig, 'etazoom', timestep) - - # plot stress field with position of trenches and ridges - if 'str' in conf.plates.plot: - fig, axis, _, _ = field.plot_scalar(step, 'sII', - vmin=0, vmax=300) - - # Annotation with time and step - axis.text(1., 0.9, str(round(time, 0)) + ' My', - transform=axis.transAxes) - axis.text(1., 0.1, str(timestep), - transform=axis.transAxes) - - # Put arrow where ridges and trenches are - plot_plate_limits_field(axis, rcmb, ridges, trenches) - - saveplot(fig, 's', timestep, close=conf.plates.zoom is None) - - # Zoom - if conf.plates.zoom is not None: - axis.set_xlim(xzoom - ladd, xzoom + radd) - axis.set_ylim(yzoom - dadd, yzoom + uadd) - saveplot(fig, 'szoom', timestep) - - # calculate stresses in the lithosphere - lithospheric_stress(step, trenches, ridges, time) - - # plotting the principal deviatoric stress field - if 'sx' in conf.plates.plot: - fig, axis, _, _ = field.plot_scalar(step, 'sII', - alpha=0.1) - - # plotting continents - cbar = conf.field.colorbar - conf.field.colorbar = False - field.plot_scalar(step, 'c', continentsfld, axis, - cmap='cool_r', vmin=0, vmax=0) - cmap2 = plt.cm.ocean - cmap2.set_over('m') - conf.field.colorbar = cbar - - # plotting principal deviatoric stress - field.plot_vec(axis, step, 'sx') - - # Annotation with time and step - axis.text(1., 0.9, str(round(time, 0)) + ' My', - transform=axis.transAxes) - axis.text(1., 0.1, str(timestep), - transform=axis.transAxes) - - # Put arrow where ridges and trenches are - plot_plate_limits_field(axis, rcmb, ridges, trenches) - - saveplot(fig, 'sx', timestep) + fig, axis, _, _ = field.plot_scalar(snap, fieldname) + + if conf.plates.continents: + c_field = np.ma.masked_where( + ~_continents_location(snap, at_surface=False), + snap.fields['c'].values[0, :, :, 0]) + cbar = conf.field.colorbar + conf.field.colorbar = False + cmap = colors.ListedColormap(["k", "g", "m"]) + field.plot_scalar(snap, 'c', c_field, axis, cmap=cmap, + norm=colors.BoundaryNorm([2, 3, 4, 5], cmap.N)) + conf.field.colorbar = cbar + + # plotting velocity vectors + field.plot_vec(axis, snap, 'sx' if conf.plates.stress else 'v') + + # Put arrow where ridges and trenches are + phi = snap.geom.p_centers + itrenches, iridges = detect_plates(snap, conf.plates.vzratio) + _plot_plate_limits_field(axis, snap.geom.rcmb, + phi[itrenches], phi[iridges]) + + saveplot(fig, f'plates_{fieldname}', snap.isnap, + close=conf.plates.zoom is None) + + # Zoom + if conf.plates.zoom is not None: + if not 0 <= conf.plates.zoom <= 360: + raise error.InvalidZoomError(conf.plates.zoom) + if 45 < conf.plates.zoom <= 135: + ladd, radd, uadd, dadd = 0.8, 0.8, 0.05, 0.1 + elif 135 < conf.plates.zoom <= 225: + ladd, radd, uadd, dadd = 0.05, 0.1, 0.8, 0.8 + elif 225 < conf.plates.zoom <= 315: + ladd, radd, uadd, dadd = 0.8, 0.8, 0.1, 0.05 + else: # >315 or <=45 + ladd, radd, uadd, dadd = 0.1, 0.05, 0.8, 0.8 + xzoom = (snap.geom.rcmb + 1) * np.cos(np.radians(conf.plates.zoom)) + yzoom = (snap.geom.rcmb + 1) * np.sin(np.radians(conf.plates.zoom)) + axis.set_xlim(xzoom - ladd, xzoom + radd) + axis.set_ylim(yzoom - dadd, yzoom + uadd) + saveplot(fig, f'plates_zoom_{fieldname}', snap.isnap) def cmd(): @@ -739,51 +328,44 @@ def cmd(): conf.core """ sdat = StagyyData() - conf.plates.plot = set_of_vars(conf.plates.plot) - if not conf.plates.vzcheck: - conf.scaling.dimensional = True - conf.scaling.factors['Pa'] = 'M' - main_plates(sdat) - else: - nb_plates = [] - time = [] - istart, iend = None, None + + isurf = _isurf(next(iter(sdat.walk))) + vrms_surf = sdat.walk.filter(rprofs=True)\ + .rprofs_averaged['vhrms'].values[isurf] + nb_plates = [] + time = [] + istart, iend = None, None + + with open(f'plates_trenches_{sdat.walk.stepstr}.dat', 'w') as fid: + fid.write('# istep time time_My phi_trench vel_trench ' + 'distance phi_cont age_trench_My\n') for step in sdat.walk.filter(fields=['T']): # could check other fields too - if conf.plates.timeprofile: + _write_trench_diagnostics(step, vrms_surf, fid) + plot_at_surface(step, conf.plates.plot) + plot_scalar_field(step, conf.plates.field) + if conf.plates.nbplates: time.append(step.timeinfo.loc['t']) - istart = step.isnap if istart is None else istart - iend = step.isnap - phi = step.geom.p_centers - limits, dvphi, vphi_surf = detect_plates_vzcheck(step) - limits.sort() - sizeplates = [phi[limits[0]] + 2 * np.pi - phi[limits[-1]]] - for lim in range(1, len(limits)): - sizeplates.append(phi[limits[lim]] - phi[limits[lim - 1]]) - phi_lim = [phi[i] for i in limits] - dvp_lim = [dvphi[i] for i in limits] - fig, axes = plt.subplots(nrows=2, sharex=True, figsize=(6.4, 6.4)) - axes[0].plot(step.geom.p_walls, vphi_surf) - axes[0].set_ylabel(r"Horizontal velocity $v_\phi$") - axes[1].plot(phi, dvphi) - axes[1].scatter(phi_lim, dvp_lim, color='red') - axes[1].set_ylabel(r"$dv_\phi/d\phi$") - axes[1].set_xlabel(r"$\phi$") - saveplot(fig, 'plate_limits', step.isnap) - fig, axis = plt.subplots() - axis.hist(sizeplates, 10, (0, np.pi)) - axis.set_ylabel("Number of plates") - axis.set_xlabel(r"$\phi$-span") - saveplot(fig, 'plate_size_distribution', step.isnap) - - nb_plates.append(len(limits)) - - if conf.plates.timeprofile: - for i in range(2, len(nb_plates) - 3): - nb_plates[i] = (nb_plates[i - 2] + nb_plates[i - 1] + - nb_plates[i] + nb_plates[i + 1] + - nb_plates[i + 2]) / 5 + itr, ird = detect_plates(step, conf.plates.vzratio) + nb_plates.append(itr.size + ird.size) + istart = step.isnap if istart is None else istart + iend = step.isnap + if conf.plates.distribution: + phi = step.geom.p_centers + itr, ird = detect_plates(step, conf.plates.vzratio) + limits = np.concatenate((itr, ird)) + limits.sort() + sizeplates = [phi[limits[0]] + 2 * np.pi - phi[limits[-1]]] + for lim in range(1, len(limits)): + sizeplates.append(phi[limits[lim]] - phi[limits[lim - 1]]) + fig, axis = plt.subplots() + axis.hist(sizeplates, 10, (0, np.pi)) + axis.set_ylabel("Number of plates") + axis.set_xlabel(r"$\phi$-span") + saveplot(fig, 'plates_size_distribution', step.isnap) + + if conf.plates.nbplates: figt, axis = plt.subplots() axis.plot(time, nb_plates) axis.set_xlabel("Time") diff --git a/tests/test_commands.py b/tests/test_commands.py index 4279a945..6ac7529a 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -17,7 +17,7 @@ def test_var_cmd(capsys): stagpy.commands.var_cmd() output = capsys.readouterr() expected = re.compile( - r'field:\n.*\nrprof:\n.*\ntime:\n.*\nplates:\n.*$', + r'field:\n.*\nrprof:\n.*\ntime:\n.*\n.*$', flags=re.DOTALL) assert expected.fullmatch(output.out)