Warning

You are reading a version of the website built against the unstable main branch. This content is liable to change without notice and may be inappropriate for your use case. You can find the documentation for the current stable release here.

firedrake.ensemble package

Submodules

firedrake.ensemble.ensemble module

class firedrake.ensemble.ensemble.Ensemble(comm, M, **kwargs)[source]

Bases: object

Create a set of space and ensemble subcommunicators.

Parameters:

  • comm – The communicator to split.

  • M – the size of the communicators used for spatial parallelism.

Keyword Arguments:

ensemble_name – string used as communicator name prefix, for debugging.

Raises:

ValueError – if M does not divide comm.size exactly.

allreduce(f, f_reduced, op=<mpi4py.MPI.Op object>)[source]

Allreduce a function f into f_reduced over ensemble_comm .

Parameters:

  • f – The a Function to allreduce.

  • f_reduced – the result of the reduction.

  • op – MPI reduction operator. Defaults to MPI.SUM.

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or if the functions are in different spaces

bcast(f, root=0)[source]

Broadcast a function f over ensemble_comm from rank root

Parameters:

  • f – The Function to broadcast.

  • root – rank to broadcast from. Defaults to 0.

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

property ensemble_rank

The rank of the local ensemble member.

property ensemble_size

The number of ensemble members.

iallreduce(f, f_reduced, op=<mpi4py.MPI.Op object>)[source]

Allreduce (non-blocking) a function f into f_reduced over ensemble_comm .

Parameters:

  • f – The a Function to allreduce.

  • f_reduced – the result of the reduction.

  • op – MPI reduction operator. Defaults to MPI.SUM.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or if the functions are in different spaces

ibcast(f, root=0)[source]

Broadcast (non-blocking) a function f over ensemble_comm from rank root

Parameters:

  • f – The Function to broadcast.

  • root – rank to broadcast from. Defaults to 0.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

irecv(f, source=-1, tag=-1)[source]

Receive (non-blocking) a function f over ensemble_comm from another ensemble rank.

Parameters:

  • f – The a Function to receive into

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • tag – the tag of the message. Defaults to MPI.ANY_TAG.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

ireduce(f, f_reduced, op=<mpi4py.MPI.Op object>, root=0)[source]

Reduce (non-blocking) a function f into f_reduced over ensemble_comm to rank root

Parameters:

  • f – The a Function to reduce.

  • f_reduced – the result of the reduction on rank root.

  • op – MPI reduction operator. Defaults to MPI.SUM.

  • root – rank to reduce to. Defaults to 0.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or is the functions are in different spaces

isend(f, dest, tag=0)[source]

Send (non-blocking) a function f over ensemble_comm to another ensemble rank.

Parameters:

  • f – The a Function to send

  • dest – the rank to send to

  • tag – the tag of the message. Defaults to 0.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

isendrecv(fsend, dest, sendtag=0, frecv=None, source=-1, recvtag=-1)[source]

Send a function fsend and receive a function frecv over ensemble_comm to another ensemble rank.

Parameters:

  • fsend – The a Function to send.

  • dest – the rank to send to.

  • sendtag – the tag of the send message. Defaults to 0.

  • frecv – The a Function to receive into.

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • recvtag – the tag of the received message. Defaults to MPI.ANY_TAG.

Returns:

list of MPI.Request objects (one for each of fsend.subfunctions and frecv.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

recv(f, source=-1, tag=-1, statuses=None)[source]

Receive (blocking) a function f over ensemble_comm from another ensemble rank.

Parameters:

  • f – The a Function to receive into

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • tag – the tag of the message. Defaults to MPI.ANY_TAG.

  • statuses – MPI.Status objects (one for each of f.subfunctions or None).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

reduce(f, f_reduced, op=<mpi4py.MPI.Op object>, root=0)[source]

Reduce a function f into f_reduced over ensemble_comm to rank root

Parameters:

  • f – The a Function to reduce.

  • f_reduced – the result of the reduction on rank root.

  • op – MPI reduction operator. Defaults to MPI.SUM.

  • root – rank to reduce to. Defaults to 0.

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or is the functions are in different spaces

send(f, dest, tag=0)[source]

Send (blocking) a function f over ensemble_comm to another ensemble rank.

Parameters:

  • f – The a Function to send

  • dest – the rank to send to

  • tag – the tag of the message. Defaults to 0

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

sendrecv(fsend, dest, sendtag=0, frecv=None, source=-1, recvtag=-1, status=None)[source]

Send (blocking) a function fsend and receive a function frecv over ensemble_comm to another ensemble rank.

Parameters:

  • fsend – The a Function to send.

  • dest – the rank to send to.

  • sendtag – the tag of the send message. Defaults to 0.

  • frecv – The a Function to receive into.

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • recvtag – the tag of the received message. Defaults to MPI.ANY_TAG.

  • status – MPI.Status object or None.

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

firedrake.ensemble.ensemble_function module

class firedrake.ensemble.ensemble_function.EnsembleCofunction(function_space: EnsembleDualSpace)[source]

Bases: EnsembleFunctionBase

A mixed finite element Cofunction distributed over an ensemble.

Parameters:

function_space (EnsembleDualSpace) – The function space of the cofunction.

class firedrake.ensemble.ensemble_function.EnsembleFunction(function_space: EnsembleFunctionSpaceBase)[source]

Bases: EnsembleFunctionBase

A mixed Function defined on a Ensemble. The subcomponents are distributed over the ensemble members, and are specified locally in an EnsembleFunctionSpace.

Parameters:

function_space (EnsembleFunctionSpace.) – The function space of the Function.

Notes

Passing an EnsembleDualSpace to EnsembleFunction will return an instance of EnsembleCofunction.

This class does not carry UFL symbolic information, unlike a Function. UFL expressions can only be defined locally on each ensemble member using a Function from EnsembleFunction.subfunctions.

See also

EnsembleFunctionSpace, EnsembleFunction, EnsembleDualSpace, EnsembleCofunction

norm(*args, **kwargs)[source]

Compute the norm of the function.

Any arguments are forwarded to norm().

class firedrake.ensemble.ensemble_function.EnsembleFunctionBase(function_space: EnsembleFunctionSpaceBase)[source]

Bases: EnsembleFunctionMixin

A mixed (co)function defined on a Ensemble. The subcomponents are distributed over the ensemble members, and are specified locally in an EnsembleFunctionSpace.

Parameters:

function_space – The function space of the (co)function.

Notes

Passing an EnsembleDualSpace to EnsembleFunction will return an instance of EnsembleCofunction.

This class does not carry UFL symbolic information, unlike a Function. UFL expressions can only be defined locally on each ensemble member using a ~firedrake.function.Function from EnsembleFunction.subfunctions.

See also

ensemble_functionspace.EnsembleFunctionSpace, ensemble_functionspace.EnsembleDualSpace, EnsembleFunction, EnsembleCofunction

assign(other, subsets=None)[source]

Set the EnsembleFunction to the value of another EnsembleFunction other.

Parameters:
copy()[source]

Return a deep copy of the EnsembleFunction.

function_space()[source]
riesz_representation(**kwargs)[source]

Return the Riesz representation of this EnsembleFunction with respect to the given Riesz map.

Internally delegates to the firedrake.function.Function.riesz_representation() of each component.

Parameters:

  • riesz_map – The Riesz map to use (l2, L2, or H1). This can also be a callable.

  • kwargs – other arguments to be passed to the firedrake.riesz_map.

property subfunctions

The (co)functions on the local ensemble member.

vec()[source]

Context manager for the global PETSc.Vec with read/write access.

It is invalid to access the Vec outside of a context manager.

vec_ro()[source]

Context manager for the global PETSc.Vec with read only access.

It is invalid to access the Vec outside of a context manager.

vec_wo()[source]

Context manager for the global PETSc.Vec with write only access.

It is invalid to access the Vec outside of a context manager.

zero(subsets=None)[source]

Set values to zero.

Parameters:

subsets (Collection[Optional[pyop2.types.set.Subset]]) – One subset for each local firedrake.function.Function. None elements will be ignored. The values of each local function will only be zeroed on the nodes on the corresponding subset.

firedrake.ensemble.ensemble_functionspace module

class firedrake.ensemble.ensemble_functionspace.EnsembleDualSpace(local_spaces: Collection, ensemble: Ensemble)[source]

Bases: EnsembleFunctionSpaceBase

A mixed dual function space defined on an ensemble.Ensemble. The subcomponents are distributed over the ensemble members, but are specified locally on each ensemble member.

Parameters:

  • local_spaces (Collection) – The list of dual function spaces on the local ensemble.comm.

  • ensemble (.ensemble.Ensemble) – The communicator that the function space is defined over.

Notes

Passing a list of dual local_spaces to EnsembleFunctionSpace will return an instance of EnsembleDualSpace.

This class does not carry UFL symbolic information, unlike a FiredrakeDualSpace. UFL expressions can only be defined locally on each ensemble member using a FiredrakeDualSpace from EnsembleDualSpace.local_spaces.

Examples

If U=CG1, V=DG0, and W=U*V, we can have the nested mixed dual space U*xV*xV*xW*xU*. This can be distributed over an ensemble.Ensemble with two ensemble members by splitting into [U*xV*]x[V*xW*xU*]. The following code creates the corresponding EnsembleDualSpace:

ensemble = Ensemble(COMM_WORLD, COMM_WORLD.size//2)
mesh = UnitIntervalMesh(8, comm=ensemble.comm)
U = FunctionSpace(mesh, "CG", 1)
V = FunctionSpace(mesh, "DG", 0)
W = U*V

if ensemble.ensemble_rank == 0:
    local_spaces = [U.dual(), V.dual()]
else:
    local_spaces = [V.dual(), W.dual(), U.dual()]

efs0 = EnsembleDualSpace(local_spaces, ensemble)
efs1 = EnsembleFunctionSpace(local_spaces, ensemble)

See also

EnsembleFunctionSpace, ensemble_function.EnsembleFunction, EnsembleDualSpace, ensemble_function.EnsembleCofunction

class firedrake.ensemble.ensemble_functionspace.EnsembleFunctionSpace(local_spaces: Collection, ensemble: Ensemble)[source]

Bases: EnsembleFunctionSpaceBase

A mixed primal function space defined on an Ensemble. The subcomponents are distributed over the ensemble members, but are specified locally on each ensemble member.

Parameters:

  • local_spaces (Collection) – The list of primal function spaces on the local Ensemble.comm.

  • ensemble (Ensemble) – The communicator that the function space is defined over.

Notes

Passing a list of dual local_spaces to EnsembleFunctionSpace will return an instance of EnsembleDualSpace.

This class does not carry UFL symbolic information, unlike a FunctionSpace. UFL expressions can only be defined locally on each ensemble member using a FunctionSpace from EnsembleFunctionSpace.local_spaces.

Examples

If U=CG1, V=DG0, and W=UxV, we can have the nested mixed space UxVxVxWxU. This can be distributed over an ensemble.Ensemble with two ensemble members by splitting into [UxV]x[VxWxU]. The following code creates the corresponding EnsembleFunctionSpace:

ensemble = Ensemble(COMM_WORLD, COMM_WORLD.size//2)
mesh = UnitIntervalMesh(8, comm=ensemble.comm)
U = FunctionSpace(mesh, "CG", 1)
V = FunctionSpace(mesh, "DG", 0)
W = U*V

if ensemble.ensemble_rank == 0:
    local_spaces = [U, V]
else:
    local_spaces = [V, W, U]

efs = EnsembleFunctionSpace(local_spaces, ensemble)

See also

EnsembleFunctionSpace, ensemble_function.EnsembleFunction, EnsembleDualSpace, ensemble_function.EnsembleCofunction

class firedrake.ensemble.ensemble_functionspace.EnsembleFunctionSpaceBase(local_spaces: Collection, ensemble: Ensemble)[source]

Bases: object

Base class for mixed function spaces defined on an Ensemble. The subcomponents are distributed over the ensemble members, and are specified locally.

Parameters:

  • local_spaces (Collection) – The list of function spaces on the local ensemble.comm.

  • ensemble (~.ensemble.Ensemble) – The communicator that the function space is defined over.

Notes

Passing a list of dual local_spaces to EnsembleFunctionSpace will return an instance of EnsembleDualSpace.

This class does not carry UFL symbolic information, unlike a FunctionSpace. UFL expressions can only be defined locally on each ensemble member using a FunctionSpace from EnsembleFunctionSpace.local_spaces.

See also

EnsembleFunctionSpace, EnsembleDualSpace, ensemble_function.EnsembleFunction, ensemble_function.EnsembleCofunction

property comm

The spatial communicator from the Ensemble communicator.

create_vec()[source]

Return a PETSc Vec on the Ensemble.global_comm with the same layout as an EnsembleFunction or EnsembleCofunction in this function space.

dual()[source]

The dual to this function space. A EnsembleDualSpace if self is a EnsembleFunctionSpace, and vice-versa.

property ensemble

The Ensemble that the function space is defined over

property ensemble_comm

The ensemble communicator from the Ensemble communicator.

property global_comm

The global communicator from the Ensemble communicator.

property global_spaces_offset

Index of the first local subspace in the global mixed space.

property layout_vec
property local_spaces

The FunctionSpace on the local ensemble.comm.

mesh()[source]

The Mesh on the local ensemble.comm.

property nglobal_dofs

The total number of dofs across all subspaces on all ensemble ranks.

property nglobal_spaces

The total number of subspaces across all ensemble ranks.

property nlocal_comm_dofs

The total number of dofs across all subspaces on the local ensemble.comm.

property nlocal_rank_dofs

The total number of dofs across all subspaces on the local MPI rank.

property nlocal_spaces

The number of subspaces on this ensemble rank.

firedrake.ensemble.ensemble_mat module

firedrake.ensemble.ensemble_mat.EnsembleBlockDiagonalMat(block_mats: Iterable, row_space: EnsembleFunctionSpaceBase, col_space: EnsembleFunctionSpaceBase)[source]

A Mat for a block diagonal matrix defined over an Ensemble. Each block acts on a single subspace of an EnsembleFunctionSpace. This is a convenience function to create a PETSc.Mat with a EnsembleBlockDiagonalMatCtx Python context.

Parameters:

  • block_mats (Iterable[PETSc.Mat]) – The PETSc Mats for each block. On each ensemble rank there must be as many Mats as there are local subspaces of row_space and col_space, and the Mat sizes must match the sizes of the corresponding subspaces.

  • row_space – The function space that the matrix acts on. Must have the same number of subspaces on each ensemble rank as col_space.

  • col_space – The function space for the result of the matrix action. Must have the same number of subspaces on each ensemble rank as row_space.

Returns:

The PETSc.Mat with an EnsembleBlockDiagonalMatCtx Python context.

Return type:

PETSc.Mat

See also

EnsembleBlockDiagonalMatCtx, EnsembleBJacobiPC

class firedrake.ensemble.ensemble_mat.EnsembleBlockDiagonalMatCtx(block_mats: Iterable, row_space: EnsembleFunctionSpaceBase, col_space: EnsembleFunctionSpaceBase)[source]

Bases: EnsembleMatCtxBase

A python Mat context for a block diagonal matrix defined over an Ensemble. Each block acts on a single subspace of an EnsembleFunctionSpace.

Parameters:

  • block_mats (Iterable[PETSc.Mat]) – The PETSc Mats for each block. On each ensemble rank there must be as many Mats as there are local subspaces of row_space and col_space, and the Mat sizes must match the sizes of the corresponding subspaces.

  • row_space – The function space that the matrix acts on. Must have the same number of subspaces on each ensemble rank as col_space.

  • col_space – The function space for the result of the matrix action. Must have the same number of subspaces on each ensemble rank as row_space.

Notes

This is a python context, not an actual PETSc.Mat. To create the corresponding PETSc.Mat users should call EnsembleBlockDiagonalMat().

See also

EnsembleBlockDiagonalMat, EnsembleBJacobiPC

mult_impl(A, x, y)[source]

Apply the action of the matrix to x, putting the result in y.

y is not guaranteed to be zero on entry. This is a convenience method allowing the matrix action to be implemented in terms of EnsembleFunction input and outputs by inheriting classes.

Parameters:

  • A (PETSc.Mat) – The PETSc matrix that self is the python context of.

  • x – The vector acted on by the matrix.

  • y – The result of the matrix action.

See also

EnsembleMatCtxBase.mult

setUp(mat)[source]
view(mat, viewer=None)[source]
class firedrake.ensemble.ensemble_mat.EnsembleMatCtxBase(row_space: EnsembleFunctionSpaceBase, col_space: EnsembleFunctionSpaceBase)[source]

Bases: object

Base class for python type Mats defined over an Ensemble.

Parameters:

  • row_space – The function space that the matrix acts on. Must have the same number of subspaces on each ensemble rank as col_space.

  • col_space – The function space for the result of the matrix action. Must have the same number of subspaces on each ensemble rank as row_space.

Notes

The main use of this base class is to enable users to implement the matrix action as acting on and resulting in an EnsembleFunction. This is done by implementing the mult_impl method.

See also

ensemble_pc.EnsemblePCBase

mult(A, x, y)[source]

Apply the action of the matrix to x, putting the result in y.

This method will be called by PETSc with x and y as Vecs, and acts as a wrapper around the mult_impl method which has x and y as EnsembleFunction for convenience. y is not guaranteed to be zero on entry.

Parameters:

  • A (PETSc.Mat) – The PETSc matrix that self is the python context of.

  • x (PETSc.Vec) – The vector acted on by the matrix.

  • y (PETSc.Vec) – The result of the matrix action.

See also

EnsembleMatCtxBase.mult_impl

mult_impl(A, x: EnsembleFunctionBase, y: EnsembleFunctionBase)[source]

Apply the action of the matrix to x, putting the result in y.

y is not guaranteed to be zero on entry. This is a convenience method allowing the matrix action to be implemented in terms of EnsembleFunction input and outputs by inheriting classes.

Parameters:

  • A (PETSc.Mat) – The PETSc matrix that self is the python context of.

  • x – The vector acted on by the matrix.

  • y – The result of the matrix action.

See also

EnsembleMatCtxBase.mult

firedrake.ensemble.ensemble_pc module

class firedrake.ensemble.ensemble_pc.EnsembleBJacobiPC[source]

Bases: EnsemblePCBase

A python PC context for a block Jacobi method defined over an Ensemble. Each block acts on a single subspace of an EnsembleFunctionSpace and is (approximately) solved with its own KSP, which defaults to -ksp_type preonly.

Available options:

  • -pc_use_amat - use Amat to apply block of operator in inner Krylov method

  • -sub_%d - set options for the %d’th block, numbered from ensemble rank 0.

  • -sub_ - set default options for all blocks.

Notes

Currently this is only implemented for EnsembleBlockDiagonalMatCtx matrices.

See also

EnsembleBlockDiagonalMatCtx, EnsembleBlockDiagonalMat

apply_impl(pc, x, y)[source]
initialize(pc)[source]

Initialize any state in this preconditioner.

This method is only called on the first time that the setUp method is called.

prefix = 'ebjacobi_'

The options prefix of this PC.

update(pc)[source]

Update any state in this preconditioner.

This method is called the on second and later times that the setUp method is called.

This method is not needed for all preconditioners and can often be a no-op.

view(pc, viewer=None)[source]

Write a basic description of this PC.

class firedrake.ensemble.ensemble_pc.EnsemblePCBase[source]

Bases: PCBase

Base class for python type PCs defined over an Ensemble.

The pc operators must be python Mats with EnsembleMatCtxBase.

Notes

The main use of this base class is to enable users to implement the preconditioner action as acting on and resulting in an EnsembleFunction. This is done by implementing the apply_impl method.

See also

EnsembleMatCtxBase

apply(pc, x, y)[source]

Apply the preconditioner to x, putting the result in y.

Both x and y are PETSc Vecs, y is not guaranteed to be zero on entry.

apply_impl(pc, x, y)[source]
initialize(pc)[source]

Initialize any state in this preconditioner.

This method is only called on the first time that the setUp method is called.

needs_python_pmat = True

Set this to False if the P matrix needs to be Python (matfree).

firedrake.ensemble.ensemble_pc.obj_name(obj)[source]

Module contents