joe_lab package

Submodules

joe_lab.initial_states module

joe_lab.initial_states.builtin_initial_state(initial_state_kw)[source]

Pulls an initial state from a catalogue of built-in options.

Parameters

initial_state_kwstr

Name of the initial state. Acceptable values: ‘sine’, ‘gaussian_even’, ‘gaussian_even_alt’, ‘gaussian_no_parity’, ‘gaussian_odd’, ‘kdv_soliton’, ‘kdv_multisoliton’, ‘gardner_soliton’, ‘bbm_solitary_wave’, ‘bbm_multisolitary’, ‘translational_mode’, ‘internal_mode’, ‘tritone’, ‘trivial’, ‘0_energy’, ‘ks_chaos’, ‘bbm_weird_wavepacket’, ‘sinegordon_soliton_interaction’, ‘sinegordon_soliton_interaction_alt’.

Returns

outinitial_state

initial_state instance representing the specified choice.

joe_lab.joe module

joe_lab.joe.do_refinement_study(model, initial_state, length, T, Ns, dts, method_kw='etdrk4', bc='periodic', sponge_params=None, show_figure=True, save_figure=False, usetex=True, dpi=400, fit_min=3, fit_max=7)[source]

Performs a refinement study based on Richardson extrapolation for error estimation, and reports the slope of “estimated error” vs. “\(\Delta t\)” on a loglog plot (computed via np.polyfit https://numpy.org/doc/stable/reference/generated/numpy.polyfit.html). Useful for quickly and painlessly validating the accuracy of an implementation.

Parameters

modelmodel

Instance of the class model.

initial_stateinitial_state

Instance of the class initial_state.

lengthfloat

Length of the spatial domain.

Tfloat

Total physical runtime of the simulation. That is, the spacetime grid covers the time interval [0,T].

Nslist of ints

Number of spatial locations on which to sample the solution to our PDE. Get an error curve for each entry in Ns.

dtslist of floats

Time step sizes for numerical integration of the PDE.

method_kwstr, optional

Name of the numerical method. Currently, ‘etdrk1’, ‘ifrk4’ are available for first-order-in-time problems, and ‘etdrk4’ is available for all problems: see timestepper. Default: ‘etdrk4’.

bcstr, optional

String specifying the boundary conditions. Must be ‘periodic’ or ‘sponge_layer’. Default: ‘periodic’.

sponge_paramsdict, optional

Contains particular parameters for the sponge layer, see damping_coeff_lt(). Default: None.

dpiint, optional

Dots-per-inch on the image. Default: 400.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figureboolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figureboolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

fit_minint, optional.

Index of dts at which we start pulling values from the Ns[-1] error curve to approximate slope of error curve. Default: 3.

fit_maxint, optional

Index of dts at which we stop pulling values from the Ns[-1] error curve to approximate slope of error curve. Default: 7.

Returns

slopefloat

Estimated slope of the “estimated error” vs. “\(\Delta t\)” line on a loglog plot.

class joe_lab.joe.initial_state(initial_state_kw: str, initial_state_func: callable)[source]

Bases: object

A class for initial states.

Having a special class for initial states instead of just representing them as callables or keywords allows for greater flexibility, as well as some symmetry with the class model.

Attributes

initial_state_kwstr

Name of the initial state.

initial_state_funcdict

A dict defined such that initial_state_func[‘initial_state_func’] is a callable representing the actual initial state function (the input initial_state_func). One must be careful of the shape of this function’s outputs when dealing with second-order problems.

initial_statendarray

The values of the initial state on our particular grid. Not defined during init, and should only be accessed under-the-hood.

class joe_lab.joe.model(model_kw: str, t_ord: int, symbol: callable, fourier_forcing: callable, nonlinear=True, complex=False)[source]

Bases: object

A class for models, which in joe specifically refers to particular PDEs.

This defines a model PDE with the name ‘model_kw’ in the Fourier-space form

\[\left(\partial_t\right)^{\texttt{t_ord}} V(k,t) + \texttt{symbol}(k) V(k,t) = \texttt{fourier_forcing}\]

where \(V(k,t)\) is the Fourier state (that is, if \(u(x,t)\) is the solution to our PDE, then \(V=\mathcal{F}u\) where \(\mathcal{F}\) is the Fourier transform).

Note how the init takes in two callables for the symbol and forcing terms, but these get converted to dicts. This is to avoid making the weird ‘no-no’ of having callable attributes. The trick comes from https://stackoverflow.com/questions/35321744/python-function-as-class-attribute-becomes-a-bound-method… we instead make the “callable attributes” dicts with callable entries. This is all under the hood in the model class, so when defining and using a model object the user doesn’t need to care.

Attributes

model_kwstr

A name for the model.

t_ordint

Temporal order of derivatives in the model. Currently only t_ord=1,2 are supported.

symboldict

Symbol of the linear, constant-coefficient part of the PDE. symbol[‘symbol’] is the callable symbol in the init, with a single parameter k, representing a Fourier-space frequency variable.

fourier_forcingdict

Forcing term in the PDE, represented in Fourier space. fourier_forcing[‘fourier_forcing’] is the callable fourier_forcing in the init, and has four arguments (V, k, x, nonlinear) where V,k,x are ndarrays and nonlinear is a boolean.

nonlinearboolean, optional

True if we include nonlinearity in the model, False otherwise. Default: True.

complexboolean, optional

False if solution to the PDE is supposed to be real, True otherwise. Default: False.

class joe_lab.joe.simulation(stgrid: dict, model: model, initial_state: initial_state, bc: str, sponge_params=None, ndump=10)[source]

Bases: object

A class for simulations.

This class is the heart of joe, and most of the scientific information joe provides (including visualizations) can be accessed via the attributes of methods of a simulation instance.

You init an instance with an stgrid (space-time grid) dict, an instance of the model class (model) to specify the PDE, and an instance of the initial_state class (initial_state). Then you can call a run simulation function on a simulation object to solve the IVP and store/save the solution, or load up the solution from an experiment you’ve already done.

Attributes

lengthfloat

Length of the spatial domain.

Tfloat

Total physical runtime of the simulation. That is, the spacetime grid covers the time interval [0,T].

Nint

Number of spatial locations on which to sample the solution to our PDE.

dtfloat

Time step size for numerical integration of the PDE.

modelmodel

An instance of the model class, see model.

model_kwstr

A name for the model.

t_ordint

Temporal order of derivatives in the model. Currently only t_ord=1,2 are supported.

initial_state_kwstr

Name of the initial state.

Udatandarray

Values of the solution to our PDE sampled on our space-time grid, stored as an array. The 0 axis represents temporal variation, and the 1 axis represents spatial variation (so each row is a fixed-time snapshot of the solution on our spatial grid).

nonlinearboolean

True if we include nonlinearity in the model, False otherwise. Default: True.

complexboolean

False if solution to the PDE is supposed to be real, True otherwise. Default: False.

sponge_paramsdict or None

Contains particular parameters for the sponge layer, see damping_coeff_lt(). Default: None.

sponge_layerboolean

True if sponge layer is included, False otherwise. Default: False.

sfracfloat

Fraction of the spatial grid that is not taken up by the sponge layer. By convention, this fraction is taken from the middle of the grid.

xndarray

Spatial grid of [-0.5*length, 0.5*length) with N points.

initial_statendarray

Initial state function sampled at x.

fmndarray

First moments (spatial integrals) of the solution to our PDE.

fm_errorndarray

Absolute difference between fm and fm[0], the initial first moment.

smndarray

Second moments (spatial \(L^2\) norms) of the solution to our PDE.

sm_errorndarray

Absolute difference between sm and sm[0], the initial second moment.

ndumpint, optional

Only every ndump timesteps are stored when saving simulation results. Default: 10.

filenamestr

Name of the .pkl file the sim may be saved in, or loaded from.

ic_picnamestr

Name of the initial state plot as a .png file.

picnamestr

Name of the Hovmoeller (filled space-time contour) plot of the solution as a .png file.

realpicnamestr

Name of the Hovmoeller plot of the real part of the solution as a .png file.

imagpicnamestr

Name of the Hovmoeller plot of the imaginary part of the solution as a .png file.

modpicnamestr

Name of the Hovmoeller plot of the modulus of the solution as a .png file.

movienamestr

Name of the movie of the solution as a .mp4 file.

realmovienamestr

Name of the movie of the real part of the solution as a .mp4 file.

imagmovienamestr

Name of the movie of the imaginary part of the solution as a .mp4 file.

modmovienamestr

Name of the movie of the modulus of the solution as a .mp4 file.

combomovienamestr

Name of the movie of the solution (or its modulus if complex=True) and its power spectrum as a .mp4 file.

get_fm()[source]

Get the first moment (spatial integral) of the solution to our PDE, fm, at each sampled time.

Also gets the absolute difference between fm and fm[0], the initial first moment.

get_sm()[source]

Get the second moment (spatial \(L^2\) norm) of the solution to our PDE, sm, at each sampled time.

Also gets the absolute difference between sm and sm[0], the initial second moment.

hov_plot(umin=None, umax=None, dpi=100, cmap='cmo.haline', fieldname='u', usetex=True, show_figure=True, save_figure=False)[source]

Create a Hovmoeller plot (filled contour plot in space-time) of the PDE solution.

If the PDE solution is complex-valued, this produces Hovmoeller plots of its real part and its imaginary part.

Parameters

uminfloat, optional.

Minimum field value to include, used to construct colorbar lower limit. Default: None.

umaxfloat, optional.

Maximum field value to include, used to construct colorbar upper limit. Default: None.

dpiint, optional

Dots-per-inch on the image. Default: 100.

cmapstr, optional

Name of the colormap to use. For a list of great colormaps see https://matplotlib.org/cmocean/. Default: ‘cmo.haline’.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

hov_plot_modulus(umin=None, umax=None, dpi=100, cmap='cmo.haline', fieldname='u', usetex=True, show_figure=True, save_figure=False)[source]

The same as hov_plot(), but only the modulus of the field is plotted.

Parameters

uminfloat, optional.

Minimum field value to include, used to construct colorbar lower limit. Default: None.

umaxfloat, optional.

Maximum field value to include, used to construct colorbar upper limit. Default: None.

dpiint, optional

Dots-per-inch on the image. Default: 100.

cmapstr

Name of the colormap to use. For a list of great colormaps see https://matplotlib.org/cmocean/. Default: ‘cmo.haline’.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

load()[source]

Loads a sim from a .pkl file that already exists in the joe_sim_archive directory.

Since time-stepping only fills the Udata attribute to the simulation instance, “loading” a saved sim just means

  1. loading the pickle and

  2. adding the Udata attribute to our simulation instance.

load_or_run(method_kw='etdrk4', print_runtime=True, save=True, verbose=True)[source]

Load our simulation from the joe_sim_archive directory if it’s there, or run the simulation if it can’t be found in joe_sim_archive. The end result is that our simulation’s Udata attribute becomes populated with the space-time grid values of our solution.

The development team recommends you use this method as your on-the-ground alternative to run_sim.

Parameters

method_kwstr, optional

Name of the numerical method. Currently, ‘etdrk1’, ‘ifrk4’ are available for first-order-in-time problems, and ‘etdrk4’ is available for all problems: see time_stepper. Default: ‘etdrk4’.

print_runtime: boolean, optional

True if you would like to print the wall-clock runtime of the simulation, False otherwise.

saveboolean, optional

True if you definitely want the simulation saved, and False otherwise. Default: True.

verboseboolean, optional

True if you want to see printed messages about whether a saved simulation was found in the joe_sim_archive directory.

plot_initial_condition(custom_ylim=None, dpi=100, color='xkcd:cerulean', fieldname='u', usetex=True, show_figure=True, save_figure=False)[source]

Produce a line plot of the initial state. If the initial state is complex-valued, plots of both the real and imaginary parts are produced.

Parameters

custom_ylimarray_like, optional.

Custom values for the y-limits on the axes. Default: None.

dpiint, optional

Dots-per-inch on the image. Default: 100.

colorstr, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:cerulean’.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all.

run_sim(method_kw='etdrk4', print_runtime=True)[source]

Perform the time-stepping starting from the initial state and ending at time T. Populates the attribute Udata (the actual values of our solution throughout the simulation) in our simulation instance.

Parameters

method_kwstr, optional

Name of the numerical method. Currently, ‘etdrk1’, ‘ifrk4’ are available for first-order-in-time problems, and ‘etdrk4’ is available for all problems: see time_stepper. Default: ‘etdrk4’.

print_runtime: boolean, optional

True if you would like to print the wall-clock runtime of the simulation, False otherwise.

save()[source]

Save the simulation object to an external .pkl file in the joe_sim_archive directory using the pickle module.

For a technical python reason, calling this function reassigns the simulation’s model attribute to None.

save_combomovie(fps=200, fieldname='u', fieldcolor='xkcd:ocean green', speccolor='xkcd:dark orange', usetex=True, dpi=100)[source]

Save a movie of the evolution of our PDE solution as well as a nested, tinier movie of its power spectrum.

If the PDE solution is complex-valued, we only show its modulus and power spectrum instead.

Parameters

fps: int, optional

Frames-per-second for the movie. Default: 200.

dpiint, optional

Dots-per-inch on the image. Default: 100.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

fieldcolor: str, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:ocean green’.

speccolor: str, optional

Name of the color you want the curve of the power spectrum to be. Default: ‘xkcd:dark orange’.

save_movie(fps=200, fieldname='u', usetex=True, fieldcolor='xkcd:ocean green', dpi=100)[source]

Save a movie of the evolution of our solution as a .mp4 file.

If the PDE solution is complex-valued this produces movies of its real part and its imaginary part.

Parameters

fps: int, optional

Frames-per-second for the movie. Default: 200.

dpiint, optional

Dots-per-inch on the image. Default: 100.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

fieldcolor: str, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:ocean green’.

save_movie_modulus(fps=200, fieldname='u', usetex=True, fieldcolor='xkcd:ocean green', dpi=100)[source]

The same as save_movie(), but with the modulus of the field plotted instead.

Parameters

fps: int, optional

Frames-per-second for the movie. Default: 200.

dpiint, optional

Dots-per-inch on the image. Default: 100.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

fieldcolor: str, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:ocean green’.

joe_lab.models module

joe_lab.models.builtin_model(model_kw, nonlinear=True)[source]

Access a given builtin model from a list of possibilities.

Parameters

model_kwstr

Name of the model to load up. Acceptable arguments: ‘bbm’, ‘gardner’, ‘gardner-bbm’, ‘kdv’, ‘ks’, ‘phi4pert’, ‘sinegordon’.

nonlinearboolean

True if we include nonlinearity in the model, False otherwise. Default: True.

Returns

outmodel

An instance of model with the given name.

joe_lab.sponge_layer module

joe_lab.sponge_layer.clip_spongeless(z, sfrac)[source]

Obtain samples of z only coming from outside the sponge layer.

Parameters

zndarray

Viewed as samples of a function on our entire spatial grid (including the sponge layer).

sfracfloat

Fraction of the spatial grid that is not taken up by the sponge layer. By convention, this fraction is taken from the middle of the grid.

Returns

outndarray

The part of z coming only from our sponge layer.

joe_lab.sponge_layer.damping_coeff_lt(x, sponge_params)[source]

Smooth damping coefficient used in Liu-Trogdon 2023 (see References below). Only applies on left side of domain.

Parameters

xfloat or ndarray

Spatial point(s) to evaluation damping coefficient at.

sponge_paramsdict

Contains particular parameters for the sponge layer: the keys are instructively named ‘l_endpt’, ‘r_endpt’, and ‘width’. For other purposes, it is useful to also populate the dict with keys ‘expdamp_freq’ (number of steps between harsh exponential damping in the sponge), ‘damping_amplitude’ (amplitude of heat-flow coefficient in the sponge layer), ‘splitting_method_kw’ (‘naive’ or ‘strang’, default to ‘naive’), and ‘spongeless_frac’ (fraction of domain that is actually “physical” and not corrupted by the sponge): these keys are required for passing the sponge_params dict to a simulation instance.

Returns

outndarray

Values of damping coefficient at the point(s) x.

References

1

Anne Liu, Thomas Trogdon, “An artificially-damped Fourier method for dispersive evolution equations”.

https://doi.org/10.1016/j.apnum.2023.05.023 .

joe_lab.sponge_layer.rayleigh_damping(V, x, sponge_params, complex=False)[source]

Rayleigh damping term for use in second-order-in-time problems involving sponge layers. Uses the Liu-Trogdon damping function damping_coeff_lt(), but sponging occurs on both sides of the domain.

Given a Fourier-space input \(V\), this function returns samples of the Rayleigh damping forcing term.

\[\mathcal{F}\left(-\beta(x)\mathcal{F}^{-1}V\right)\]

where \(\mathcal{F}\) denotes the Fourier transform and \(\beta(x)\) denotes a damping coefficient close to 1 near the boundary of our domain and effectively zero everywhere else.

Parameters

Vcomplex ndarray

Fourier-space representation of a given function sampled at some number of points.

xndarray

Points in physical space where function is sampled.

sponge_paramsdict

Parameters of our sponge layer, see damping_coeff_lt().

complexboolean, optional.

True if the inverse FFT of V is complex and False if it is real. Default: False.

Returns

outcomplex ndarray

Fourier-space representation of the Rayleigh damping forcing term.

joe_lab.time_stepper module

joe_lab.time_stepper.assemble_damping_mat(N, length, x, dt, sponge_params, complex=False)[source]

Get the matrix that is to be inverted at each time-step in the artificial damping stage. Only important when the sponge layer is turned on.

Strictly speaking we store the matrix as a LinearOperator (https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html) for performance reasons: the array is dense and too large for direct storage and direct multiplication to be performant.

Parameters

Nint

Number of grid points (or frequencies) we sample in our simulation.

lengthfloat

Length of the spatial domain.

xndarray

Spatial grid of [-0.5*length, 0.5*length].

dtfloat

time-step size for numerical integration of the PDE.

sponge_paramsdict

Contains particular parameters for the sponge layer, see damping_coeff_lt().

complexboolean

False if solution to the PDE is supposed to be real, True otherwise. Default: False.

Returns

outscipy.sparse.linalg.LinearOperator

The damping matrix, represented by its shape and matrix-vector multiplication map.

joe_lab.time_stepper.do_diffusion_step(q, B)[source]

Solve the SPD linear system Bz = q via the conjugate gradient method. This practical shows up when doing a step of diffusion when sponge layers are included.

Parameters

qndarray

Right-hand side of the relevant linear system.

Bscipy.sparse.linalg.LinearOpeartor

Matrix of the linear system (representing a Fourier-space discretization of the diffusion operator that activates only in the sponge layer).

Returns

outndarray

Solution to the linear system.

joe_lab.time_stepper.do_etdrk1_step(V, propagator, forcing, Q1)[source]

Performs a single step of ETDRK1 (the exponential Euler method).

Parameters

Vndarray

Solution value at initial time.

propagatorndarray

Fourier-space representation of the propagator associated to the linear, constant-coefficient part of the PDE. The time-step is already encoded in here (we propagate for a length of time equal to the time-step).

forcingcallable

Fourier-space forcing term in the PDE.

Q1ndarray

Exponential Runge-Kutta weight (see get_Q1()).

Returns

outndarray

Value of solution at time = initial time + time-step length.

joe_lab.time_stepper.do_etdrk4_step(V, propagator, propagator2, forcing, greeks)[source]

Performs a single step of ETDRK4 (fourth-order exponential Runge-Kutta) for a first-order-in-time problem.

Parameters

Vndarray

Solution value at initial time.

propagatorndarray

Fourier-space representation of the propagator associated to the linear, constant-coefficient part of the PDE. The time-step is already encoded in here (we propagate for a length of time equal to the time-step).

propagator2ndarray

The same as above, but for only half the time-step.

forcingcallable

Fourier-space forcing term in the PDE.

greekslist

Entries are the Greeks/fourth-order exponential Runge-Kutta weights (see get_greeks_first_order()).

Returns

outndarray

Value of solution at time = initial time + time-step length.

joe_lab.time_stepper.do_etdrk4_step_second_order(V, propagator, propagator2, forcing, greeks)[source]

Performs a single step of ETDRK4 (fourth-order exponential Runge-Kutta) for a second-order-in-time problem.

This is different from the analogous function for first-order problems do_etdrk4_step() because we must account for the Greeks now being sparse matrices instead of simple vector-shaped arrays.

Parameters

Vndarray

Solution value at initial time.

propagatorndarray

Fourier-space representation of the propagator associated to the linear, constant-coefficient part of the PDE. the time-step is already encoded in here (we propagate for a length of time equal to the time-step).

propagator2ndarray

The same as above, but for only half the time-step.

forcingcallable

Fourier-space forcing term in the PDE.

greekslist

Entries are the Greeks/fourth-order exponential Runge-Kutta weights (see get_greeks_second_order()).

Returns

outndarray

Value of solution at time = initial time + time-step length.

joe_lab.time_stepper.do_ifrk4_step(V, propagator, propagator2, forcing, dt)[source]

Parameters

Vndarray

Solution value at initial time.

propagatorndarray

Fourier-space representation of the propagator associated to the linear, constant-coefficient part of the PDE. the time-step is already encoded in here (we propagate for a length of time equal to the time-step).

propagator2ndarray

The same as above, but for only half the time-step.

forcingcallable

Fourier-space forcing term in the PDE.

dtfloat

Time-step.

Returns

outndarray

Value of solution at time = initial time + time-step length.

joe_lab.time_stepper.do_time_stepping(sim, method_kw='etdrk4')[source]

Perform all the time-steps in the simulation, starting from time = 0 and ending at time = T.

Parameters

simsimulation

The instance of simulation our timestepper is supposed to work on.

method_kwstr

The name of the time-stepper (for instance, ‘etdrk4’).

Returns

Udatandarray

Values of the solution to our PDE sampled on our space-time grid, stored as an array. The 0 axis represents temporal variation, and the 1 axis represents spatial variation (so each row is a fixed-time snapshot of the solution on our spatial grid).

joe_lab.time_stepper.get_Q1(N, dt, z)[source]

Computes the single weight for first-order exponential quadrature for a first-order-in-time problem.

This by done by adapting the approach of Kassam and Trefethen 2005.

Parameters

Nint

Number of frequencies we sample.

dtfloat

Time-step for the simulation.

zndarray

z should be vectorial (it has shape Nx1). In practice, z is the diagonal vector of the matrix of the linear, constant-coefficient part of the PDE in Fourier space.

Returns

outndarray

Exponential Runge-Kutta weight, Q1, for the exponential Euler method.

References

1

Aly-Khan Kassam, Lloyd N. Trefethen “Fourth-Order Time-Stepping for Stiff PDEs”.

https://doi.org/10.1137/S1064827502410633 .

joe_lab.time_stepper.get_greeks_first_order(N, dt, z)[source]

Computes all the Greeks (Runge-Kutta weights for exponential quadrature) for ETDRK4 applied to a first-order-in-time problem.

This by done Pythonizing the code from Kassam and Trefethen 2005, that is, by numerically approximating Cauchy integrals.

Parameters

Nint

Number of frequencies we sample.

dtfloat

Time-step for the simulation.

zndarray

z should be vectorial (it has shape Nx1). In practice, z is the diagonal vector of the matrix of the linear, constant-coefficient part of the PDE in Fourier space.

Returns

[Q, f1, f2, f3]list

A list whose entries are the Greeks (vector-shaped arrays).

References

1

Aly-Khan Kassam, Lloyd N. Trefethen “Fourth-Order Time-Stepping for Stiff PDEs”.

https://doi.org/10.1137/S1064827502410633.

joe_lab.time_stepper.get_greeks_second_order(N, dt, A)[source]

Computes all the Greeks (Runge-Kutta weights for exponential quadrature) for EDTRK4 applied to a second-order-in-time problem.

This by done Pythonizing the code from Kassam and Trefethen 2005, that is, by numerically approximating Cauchy integrals.

Note that, since the linear, constant-coefficient part of the PDE system is no longer diagonal when we have a second-order-in-time problem, the code below is very different from the analogous function for first-order problems get_greeks_first_order(). In particular, the radius of our Cauchy integral contour must be chosen based on the argument A.

Parameters

Nint

Number of frequencies we sample.

dtfloat

Time-step for the simulation.

Andarray

A is the matrix of the linear, constant-coefficient part of the PDE in Fourier space.

Returns

[Q, f1, f2, f3]list

A list whose entries are the Greeks (arrays).

References

1

Aly-Khan Kassam, Lloyd N. Trefethen “Fourth-Order Time-Stepping for Stiff PDEs”.

https://doi.org/10.1137/S1064827502410633 .

class joe_lab.time_stepper.timestepper(method_kw, sim, scale=1.0, complex=False)[source]

Bases: object

Completely encodes a time-stepping method: instances of this class perform time-stepping during simulations.

Attributes

method_kwstr

The name of the time-stepper (for instance,’etdrk4’).

simsimulation

The instance of simulation our timestepper is supposed to work on.

t_ordint

Temporal order of derivatives in the model. Currently only t_ord=1,2 are supported.

auxlist or ndarray

Auxiliary quantities (typically Greeks/exponential Runge-Kutta weights and propagators) that are needed for time-stepping. These are usually computed in a pre-processing stage and saved in the directory joe_timestepper_aux after the instance is created but before time-stepping occurs.

auxfilename: str

Name of the .pkl file containing aux.

complexboolean

False if solution to the PDE is supposed to be real, True otherwise. Default: False.

scalefloat, optional.

An often-unused parameter related to the splitting method in the sponge layer diffusion solve. scale = 0.5 if we use Strang splitting, and scale=1 otherwise.

do_time_step(V, forcing)[source]

Do a single time-step with the method of choice.

Parameters

Vndarray

Solution value at initial time.

forcingcallable

Fourier-space forcing term in the PDE.

Returns

outndarray

Value of solution at time = initial time + time-step length.

get_aux()[source]

Obtains auxiliary timestepping quantities aux via quadrature.

load_aux()[source]

Load the auxiliary timestepping quantities aux from the directory joe_timestepper_aux.

save_aux()[source]

Save the auxiliary timestepping quantities aux in the directory joe_timestepper_aux.

joe_lab.utils module

joe_lab.utils.dealiased_pow(V, p)[source]

Compute the Fourier-space version of a power of an array, dealiased by zero-padding.

Important: this only works for real-valued arrays because for complex arrays \(u\), algebraic nonlinearities typically involve the complex conjugate \(\overline{u}\) as well as \(u\). So, padding for such nonlinearities terms should be done on-the-fly.

A future release of joe may include support for dealiased complex powers.

Parameters

V : complex ndarray

p : int

Returns

out: complex ndarray

A version of \(\mathcal{F}\left[\left(\mathcal{F}^{-1}V\right)^p\right]\) that has been dealiased via zero-padding.

joe_lab.utils.integrate(u, length)[source]

Integrates N samples of a real or complex space-time field over the spatial interval [-0.5*length, 0.5*length] using the FFT.

Parameters

undarray

Real or complex array to be integrated in space. Convention is that spatial variation is stored in the -1 axis.

lengthfloat

Total length of spatial domain.

Returns

out: float or ndarray

Approximation of the integral of the sampled space-time field over the spatial interval [-0.5*length, 0.5*length].

joe_lab.utils.my_fft(u, n=None, complex=False)[source]

Takes FFT of a real or complex array along the -1 axis, optimizing storage if the array is real.

Wraps the scipy.fft routines fft, rfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).

Parameters

undarray

Real or complex numpy array.

n: int or None, optional.

Total number of frequencies to include, if padding. Default: None (no padding).

complexboolean, optional

Equal to True if u is complex and False if u is real. Default: False.

Returns

outcomplex ndarray

Same shape conventions as scipy.fft.fft and scipy.fft.rfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).

joe_lab.utils.my_ifft(V, n=None, complex=False)[source]

Takes inverse FFT of an array along the -1 axis, optimizing storage if the inverse FFT is expected to be real.

Wraps the scipy.fft routines ifft, irfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).

Parameters

Vcomplex ndarray

Array we want to apply inverse FFT to.

n: int or None, optional.

Total number of frequencies to pad V with, if padding is desired. Default: None (no padding).

complexboolean, optional

Equal to True if inverse FFT of V is known to be complex and False if inverse FFT is known to be real. Default: False.

Returns

out: ndarray

Same shape conventions as scipy.fft.ifft and scipy.fft.irfft (see https://docs.scipy.org/doc/scipy/reference/fft.html).

joe_lab.visualization module

joe_lab.visualization.hov_plot(x, t, u, fieldname, umin=None, umax=None, dpi=100, show_figure=True, save_figure=False, picname='', cmap=<matplotlib.colors.LinearSegmentedColormap object>, usetex=True)[source]

Create a Hovmoeller plot (filled contour plot in space-time) of the space-time field u.

This just wraps matplotlib code with a bunch of nice presets, labels, etc.!

Parameters

xndarray

Spatial sample locations of the field (horizontal axis).

tndarray

Temporal sample locations of the field (vertical axis).

undarray

Samples values of the field (axis 0 is temporal variation, axis 1 is spatial variation).

uminfloat, optional.

Minimum field value to include, used to construct colorbar lower limit. Default: None.

umaxfloat, optional.

Maximum field value to include, used to construct colorbar upper limit. Default: None.

dpiint, optional

Dots-per-inch on the image. Default: 100.

cmapstr, optional

Name of the colormap to use. For a list of great colormaps see https://matplotlib.org/cmocean/. Default: ‘cmo.haline’.

fieldnamestr, optional

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

picnamestr, optional

Name of the .png file to save the figure to. Default: “”.

joe_lab.visualization.nice_hist(data, xlabel, dpi=100, show_figure=True, save_figure=False, picname='', color='xkcd:deep pink', usetex=True)[source]

Renders a nice histogram of the given data.

This just wraps matplotlib code with a bunch of nice presets, labels, etc.!

Parameters

datandarray

Data to be binned and plotted.

xlabelstr

Label of the horizontal axis.

dpiint, optional

Dots-per-inch on the image. Default: 100.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

color: str, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:deep pink’.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

picnamestr, optional

Name of the .png file to save the figure to. Default: “”.

joe_lab.visualization.nice_multiplot(xs, ys, xlabel, ylabel, curvelabels, linestyles, colors, linewidths, custom_ylim=None, dpi=100, show_figure=True, save_figure=False, picname='', usetex=True)[source]

Renders a nice 2D line plot of a bunch of y’s as a function of a bunch of x’s.

This just wraps matplotlib code with a bunch of nice presets, labels, etc.!

Parameters

xslist of ndarrays

Manipulated variable/horizontal axis.

yslist of ndarrays

Responding variable/vertical axis.

xlabelstr

Label of the horizontal axis.

ylabelstr

Label of the vertical axis.

curvelabels: list of strings

Each entry is a label for the corresponding curve, used to create the legend.

linestyleslist of strings

Style of the plotted lines, see matplotlib docs for a list of acceptable values: https://matplotlib.org/stable/api/_as_gen/matplotlib.lines.Line2D.html#matplotlib.lines.Line2D.set_linestyle

colors: list of strings

Name of the colors you want each curve to be. For a list of great colors see https://xkcd.com/color/rgb/.

linewidthslist of floats

Widths of the plotted lines.

dpiint, optional

Dots-per-inch on the image. Default: 100.

custom_ylimtuple, optional

Custom values for the limits on the vertical axis.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

picnamestr, optional

Name of the .png file to save the figure to. Default: “”.

joe_lab.visualization.nice_plot(x, y, xlabel, ylabel, dpi=100, custom_ylim=None, show_figure=True, save_figure=False, picname='', linestyle='solid', color='xkcd:cerulean', usetex=True)[source]

Renders a nice 2D line plot of y as a function of x.

This just wraps matplotlib code with a bunch of nice presets, labels, etc.!

Parameters

xndarray

Manipulated variable/horizontal axis.

yndarray

Responding variable/vertical axis.

xlabelstr

Label of the horizontal axis.

ylabelstr

Label of the vertical axis.

dpiint, optional

Dots-per-inch on the image. Default: 100.

custom_ylimtuple, optional

Custom values for the limits on the vertical axis.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

color: str, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:cerulean’.

linestylestr, optional

Style of the line, see matplotlib docs for a list of acceptable values: https://matplotlib.org/stable/api/_as_gen/matplotlib.lines.Line2D.html#matplotlib.lines.Line2D.set_linestyle .

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

picnamestr, optional

Name of the .png file to save the figure to. Default: “”.

joe_lab.visualization.plot_refinement_study(model, initial_state, length, T, Ns, dts, errors, method_kw='etdrk4', bc='periodic', show_figure=True, save_figure=False, usetex=True, dpi=400)[source]

Plots a refinement study (see do_refinement_study().

The actual code is a bit ugly but it gets the job done and prevents the main joe scripts from being too cluttered.

Parameters

model: model

Instance of the class model.

initial_state: initial_state

Instance of the class initial_state.

lengthfloat

Length of the spatial domain

Tfloat

Total physical runtime of the simulation. That is, the spacetime grid covers the time interval [0,T].

Nslist of ints

Number of spatial locations on which to sample the solution to our PDE.

dtslist of floats

Time step sizes for numerical integration of the PDE.

errorslist of ndarrays

Estimated errors for each choice of space-time grid.

method_kwstr, optional

Name of the numerical method. Currently, ‘etdrk1’, ‘ifrk4’ are available for first-order-in-time problems, and ‘etdrk4’ is available for all problems: see timestepper. Default: ‘etdrk4’.

bcstr, optional

String specifying the boundary conditions. Must be ‘periodic’ or ‘sponge_layer’.

dpiint, optional

Dots-per-inch on the image. Default: 400.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

show_figure: boolean, optional

True if you want the figure to appear in a pop-up window once it is rendered, False otherwise. Default: True.

save_figure: boolean, optional

True if you want the figure to be saved as a .png file, False if you do not want the figure to be saved at all. Default: False.

joe_lab.visualization.save_combomovie(u, x, length, dt, ndump, filename, fieldname, fps=200, periodic=True, usetex=True, fieldcolor='xkcd:ocean green', speccolor='xkcd: dark magenta', dpi=100)[source]

Save a movie of the evolution of a real scalar field as a .mp4 file, and a tinier embedded movie of its power spectrum.

This just wraps matplotlib code with a bunch of nice presets, labels, etc.!

Parameters

undarray

Sampled values of our scalar field on a given space-time grid.

lengthfloat

Length of the spatial domain.

xndarray

Spatial grid of [-0.5*length, 0.5*length].

dtfloat

Lattice spacing on the time axis.

fieldnamestr

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

ndumpint

Only every ndump timesteps are stored in u.

filenamestr

Name of the .mp4 file to save the movie as.

periodicboolean, optional

True if the field physically obeys periodic boundary conditions, False otherwise.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

fieldcolor: str, optional

Name of the color you want the physical-space curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:ocean green’.

speccolor: str, optional

Name of the color you want the Fourier-space power spectrum curve to be. Default: ‘xkcd:dark magenta’.

fps: int, optional

Frames-per-second for the movie. Default: 200.

dpiint, optional

Dots-per-inch on the image. Default: 100.

joe_lab.visualization.save_movie(u, x, length, dt, fieldname, ndump, filename, fps=200, periodic=True, usetex=True, fieldcolor='xkcd:ocean green', dpi=100)[source]

Save a movie of the evolution of a real scalar field as a .mp4 file.

This just wraps matplotlib code with a bunch of nice presets, labels, etc.!

Parameters

undarray

Sampled values of our scalar field on a given space-time grid.

lengthfloat

Length of the spatial domain.

xndarray

Spatial grid of [-0.5*length, 0.5*length].

dtfloat

Lattice spacing on the time axis.

fieldnamestr

Name of the PDE solution, which will also be the label on the picture’s y-axis. Default: ‘u’.

ndumpint

Only every ndump timesteps are stored in u.

filenamestr

Name of the .mp4 file to save the movie as.

periodicboolean, optional

True if the field physically obeys periodic boundary conditions, False otherwise.

usetexboolean, optional

True if you want to render the plot labels in TeX, and False if you want no TeX. Default: True.

fieldcolor: str, optional

Name of the color you want the curve to be. For a list of great colors see https://xkcd.com/color/rgb/. Default: ‘xkcd:ocean green’.

fps: int, optional

Frames-per-second for the movie. Default: 200.

dpiint, optional

Dots-per-inch on the image. Default: 100.

joe_lab.visualization.spinner(title: Optional[str] = None) ContextManager[source]

Displays a nice little progress bar to indicate that certain tasks are running.

The code is taken right from the reference, with no real input from the joe development team.

Parameters

title : str, optional

References

1

<https://stackoverflow.com/questions/44851940/python-cli-progress-bar-spinner-without-iteration>