def extract_coefficients(equation: sympy.Expr, local_map: dict,
global_coords: list) -> tuple:
"""
Args:
equation: The equation in local coordinates.
local_map: The mapping from local coordinates to the index of a
global coordinate.
global_coords: The list of global co-ordinates.
Returns:
The linear and nonlinear parts of the equation in the global
co-ordinate system.
Extracts the coordinates from the given equation and maps them into
the global coordinate space.
Equations are assumed to come in as sympy expressions of the form
:math:`\\Phi(x) = 0`.
local_map is a dictionary mappings
.. math::
M: \\rightarrow i
where :math:`x` are the local co-ordinates and the keys of local_map, and
the values are the indices :math:`i` such that `global_coord[i]` is the
corresponding global coordinate. The result is :math:`L,N` such that:
.. math::
Ly + N(y) = 0
"""
coefficients = {}
nonlinear_terms = sympy.S(0)
subs = [(k, global_coords[v]) for k, v in local_map.items()]
local_variables = set(local_map.keys())
subs.sort(key=lambda x: str(x[1])[-1], reverse=True)
logger.debug("Extracting coefficients from %s", repr(equation))
logger.debug("Using local-to-global substitutions %s", repr(subs))
remainder = equation
mappings = list(local_map.items())
while mappings and remainder:
variable, index = mappings.pop()
coefficient = equation.coeff(variable)
if not coefficient:
continue
remainder -= coefficient * variable
if coefficient.atoms() & local_variables:
nonlinear_terms += (coefficient * variable).subs(subs)
else:
coefficients[index] = coefficient
nonlinear_terms = sympy.expand(nonlinear_terms + remainder.subs(subs))
logger.debug("Linear terms: %s", repr(coefficients))
logger.debug("Nonlinear terms: %s", repr(nonlinear_terms))
return coefficients, nonlinear_terms
