import numpy as np
import scipy.sparse as sp
[docs]def log_history(f):
"""Creates a log for the input of a given function.
Parameters
----------
f : Callable
Function to use
Returns
-------
Callable
Wrapped function
"""
history = []
def wrapped(*args, **kwargs):
history.append((args, kwargs))
return f(*args, **kwargs)
return wrapped, history
[docs]def attribute_helper(ifs, cell_type, name, boundary=[]):
"""Helper function for flatten wrappers extracting needed
information about IndexedFaceSet attributes.
Parameters
----------
ifs : IndexedFaceSet
cell_type : str
Cell type of the attribute
name : str
Name of the indexed faceset attribute to use
boundary : Iterable
Indices of cells to ignore
Returns
-------
array : numpy.ndarray
Array of the attribute
mask : numpy.ndarray
Mask for non-boundary cells
shape_ign : tuple of ints
Shape of attribute array used for flatten wrappers, when
boundary values get ignored
shape_rem : tuple of ints
Shape of attribute array used for flatten wrappers, when
boundary values get removed
"""
c, attribute_string = ifs._return_cells(cell_type)
attr_dict = getattr(ifs, attribute_string)
# Mask for applying boundary condition
mask = np.ones(len(c), dtype=bool)
if np.size(boundary):
mask[boundary] = False
array = attr_dict[name]
shape_ign = np.shape(array)
shape_rem = np.shape(array[mask])
return (array, mask, shape_ign, shape_rem)
[docs]def nflat_index(idx, n):
"""Return indices for flattened (m,n)-dimensional arrays.
The returned indices correspond to the indices of the entries
of the rows given by idx after flattening the array.
Parameters
----------
idx : Iterable
Indices to convert to flattened indices
n : int
Dimension of target (m,n) array
Returns
-------
numpy.ndarray
Flat indices corresponding to idx
"""
return np.add.outer(n * np.array(idx), np.arange(n)).flatten()
[docs]def zero_columns(f, columns):
"""Wraps sparse matrix valued function to zero columns.
Parameters
----------
f: Callable
Function to wrap
columns: Iterable
Columns to zero
Returns
-------
Callable
wrapped function
"""
def wrapped(*args, **kwargs):
temp = f(*args, **kwargs)
n = np.ones(temp.T.shape[0])
n[columns] = 0.0
mat = sp.diags(n) # uses scipy.sparse.diags to not break type
return temp @ mat
return wrapped
[docs]def remove_columns(f, columns):
"""Wraps sparse matrix valued function to remove columns.
Parameters
----------
f: Callable
Function to wrap
columns: Iterable
Columns to remove
Returns
-------
Callable
wrapped function
"""
columns = np.unique(columns)
def wrapped(*args, **kwargs):
temp = f(*args, **kwargs)
coo = temp.tocoo() # compatible with sparse matrix formats
keep = ~np.in1d(coo.col, columns)
coo.data, coo.row, coo.col = coo.data[keep], coo.row[keep], coo.col[keep]
coo.col -= columns.searchsorted(coo.col)
coo._shape = (coo.shape[0], coo.shape[1] - len(columns))
return coo.tocsr()
return wrapped
[docs]def initial_guess(ifs, cell_type, name, boundary=[]):
"""Extracts attribute values for non-boundary cells.
Parameters
----------
ifs : IndexedFaceSet
cell_type : str
Cell type of the attribute
name : str
Name of the indexed faceset attribute to use
boundary : Iterable
Indices of cells to ignore
Returns
-------
numpy.ndarray of shape (m,)
Values of attribute [name] on non-boundary cells
in a flattened array
"""
c, attribute_string = ifs._return_cells(cell_type)
attr_dict = getattr(ifs, attribute_string)
mask = np.ones(len(c), dtype=bool)
if np.size(boundary):
mask[boundary] = False
return attr_dict[name][mask].flatten()