tfc.utils package

Submodules

tfc.utils.CeSolver module

class tfc.utils.CeSolver.CeSolver(C: list[Callable[[AppliedUndef | Expr | Any], AppliedUndef | Any]] | tuple[Callable[[AppliedUndef | Expr | Any], AppliedUndef | Any], ...], kappa: list[Expr | Any] | tuple[Expr | Any, ...], s: list[Expr | Any] | tuple[Expr | Any, ...], g: AppliedUndef | Any)

Bases: object

Constrained expression solver.

This class solves constrained expressions for you.

Parameters:

• C (ConstraintOperators) – This is a tuple or list constraint operators. Each element in the iterable should be a Python function that takes in a sympy function, such as g(x), and outputs that function evaluated in the same way as the function in the constraint. For example, if the constraint was u(3) = 0, then the assocaited constraint operator would be: lambda u: = u.subs(x,3).

• kappa (Exprs) – This is a tuple or list of the kappa portion of each constraint. For the example u(3) = 0, the kappa portion is simply 0. Note, the standard Python 0 should not be used, but rather, sympy.re(0).

• s (Exprs) – This is a tuple or list of the support functions. These should be given in terms of sympy symbols and constants. For example, if we wanted to use the constant function x = 1 as a support function, then we would use sympy.re(1) in this iterable.

• g (Union[AppliedUndef, Any]) – This is the free function used in the constrained expression. For example, g(x).

References

The algorithm used here is given at 26:13 of this video: https://www.youtube.com/watch?v=uisOZVBHA2U&t=1573s

Examples

Consider the constraints u(0) = 2 and u_x(2) = 1 where u_x is the derivative of u(x) with respect to x. Moreover, suppose we want to use g(x) as the free function.

import sympy as sp from tfc.utils import CeSolver

x = sp.Symbol(“x”) y = sp.Symbol(“y”) u = sp.Function(“u”) g = sp.Function(“g”)

C = [lambda u: u.subs(x,0), lambda u: sp.diff(u,x).subs(x,2)] K = [sp.re(2), sp.re(1)] s = [sp.re(1), x]

cs = CeSolver(C,K,s, g(x)) ce = cs.ce

In the above code example, ce is the constrained expression that satisfies these constraints.

property C: list[Callable[[AppliedUndef | Expr | Any], AppliedUndef | Any]] | tuple[Callable[[AppliedUndef | Expr | Any], AppliedUndef | Any], ...]

Constraint operators.

Returns:

Constraint operators.

Return type:

ConstraintOperators

property S: MutableDenseMatrix

Support matrix.

Returns sp.Matrix

Support matrix. The elements are on the field over which the constrained expression is defined.

property alpha: MutableDenseMatrix

Alpha matrix (inverse of the support matrix)

Returns sp.Matrix

alpha matrix. The elements are on the field over which the constrained expression is defined.

property ce: Any

Constrained expression.

Returns:

Constrained expression.

Return type:

Any

checkCe() → bool

Checks the constrained expression stored in the class against the stored constraints.

Returns:

Returns True if the constraint expression satisfies the constraints and false otherwise.

Return type:

bool

property g: AppliedUndef | Any

Free function.

Returns:

Free function.

Return type:

Union[AppliedUndef, Any]

property kappa

Kappa values.

Returns:

Kappa values.

Return type:

Exprs

property phi: Any

Switching functions.

Returns:

Switching functions.

Return type:

Any

property print_type: Literal['tfc', 'pretty', 'latex', 'str']

property rho: Any

Projection functionals.

Returns:

Projection functionals.

Return type:

Any

property s: list[Expr | Any] | tuple[Expr | Any, ...]

Switching functions.

Returns:

Support functions.

Return type:

Exprs

class tfc.utils.CeSolver.TfcPrinter(settings=None)

Bases: PythonCodePrinter

tfc.utils.Html module

class tfc.utils.Html.Dot(outFile: str, name: str)

Bases: object

This class contains helper functions used to create dot graphs.

Render(formats: list[str] = ['cmapx', 'svg'])

This function renders the dot graph as a .svg and as a .cmapx.

Parameters:

formats (List[str], optional) – List whose elementts dictate which formats to render the dot graph in. (Default value = [“cmapx”, “svg”])

class tfc.utils.Html.HTML(outFile: str)

Bases: object

This contains helper functions for creating HTML files with yattag.

Parameters:

outFile (str) – Output file

GenerateHtml() → str

This function generates and formats the HTML file text.

Returns:

htmlstr

HTML file as a string.

ReadFile(inFile: str) → str

This function reads the file specified by “inFile” and retuns the contents as a string.

Parameters:

inFile (Path) – File to read.

Returns:

outStr – Contents of inFile as a string.

Return type:

str

WriteFile()

This function writes the HTML file text to a file.

tfc.utils.Latex module

class tfc.utils.Latex.table

Bases: object

This class is used to create the text needed for latex tables.

static SimpleTable(arrIn: ndarray[Any, dtype[_ScalarType_co]], form: str = '%.4E', colHeader: list[str] | None = None, rowHeader: list[str] | None = None) → str

This function creates a simple latex table for the 2D numpy array arrIn. The “form” argument specifies the number format to be used in the tabular environment. The “colHeader” arugment is a list of strings that are used as the first row in the tabular environment. The “rowHeader” argument is a list of strings that are used in the first column of each row in the tabular environment. The latex tabular environment is returned as a string.

Parameters:

• arrIn (NDArray) – Array to convert to Latex table.

• form (str, optional) – Format string for the table numbers. (Default value = “%.4E”)

• colHeader (Optional[List[str]]) – List of strings that form the column headers. (Default value = None)

• rowHeader (Optional[List[str]]) – List of strings to use as the row headers. (Default value = None)

Returns:

table – Latex table as a string.

Return type:

str

tfc.utils.MakePlot module

class tfc.utils.MakePlot.MakePlot(xlabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str], ylabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str], titles: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str] | None = None, twinYlabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str] | None = None, zlabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str] | None = None, style: str | dict | list[str] | list[dict] | None = None)

Bases: object

This class is used to easily create journal-article-ready plots and subplots. The class can create 2D as well as 3D plots and even has support for twin y-axes.

FullScreen()

This function makes the plot fullscreen.

PartScreen(width: float, height: float)

This function makes the plot width x height inches.

Parameters:

• width (float) – Width of the plot in inches.

• height (float) – Height of the plot in inches.

animate(animFunc: Callable[[], Generator[None, None, None]], outDir: str = 'MyMovie', fileName: str = 'images', save: bool = True, delay: int = 10)

Creates an animation using a Python generator.

Parameters:

• animFunc (Callable[[], Generator[None, None, None]]) – Function whose returned generator modifies what is displayed on the plot

• outDir (Path, optional) – Directory to save frames in: only used if save = True. (Default value = “MyMovie”)

• fileName (str, optional) – Filename for the frames: only used if save = True. (Default value = “images”)

• save (bool, optional) – Controls whether the function saves the frames of the animation or not. (Default value = True)

• delay (pint, optional) – Delay in milliseconds between frames: only used if save = False. (Default value = 10)

draw()

This function draws the canvas.

save(fileName: str, transparent: bool = True, fileType: Literal['png', 'pdf', 'ps', 'eps', 'svg', None] | None = None)

This function crops and saves the figure.

Parameters:

• fileName (Path) – Filename where the figure should be saved.

• transparent (bool, optional) – Whether to save the plot with transparency or not. (Default value = True)

• fileType (Literal["png", "pdf", "ps", "eps", "svg", None], optional) – File suffix to use. If None, then the suffix will inferred from the suffix of the fileName. (Default value = None)

saveAll(fileName: str, transparent: bool = True, fileType: Literal['png', 'pdf', 'ps', 'eps', 'svg', None] | None = None)

This function invokes the save and savePickle functions.

Parameters:

• fileName (Path) – Filename where the figure should be saved.

• transparent (bool, optional) – Whether to save the plot with transparency or not. (Default value = True)

• fileType (Literal["png", "pdf", "ps", "eps", "svg", None], optional) – File suffix to use. If None, then the suffix will be inferred from the suffix of the fileName. (Default value = None)

savePickle(fileName: str)

This function saves the figure in a pickle format so it can be opened and modified later.

Parameters:

fileName (Path) – Filename where the figure should be saved. Note, this should not include the file suffix.

show()

This function shows the plot.

tfc.utils.MayaviMakePlot module

class tfc.utils.MayaviMakePlot.MakePlot

Bases: object

MakePlot class for Mayavi.

static ColormapGradient(c1: str | tuple[float, float, float, float] | ndarray[Any, dtype[float64]], c2: str | tuple[float, float, float, float] | ndarray[Any, dtype[float64]]) → ndarray[Any, dtype[uint8]]

Returns a Mayavi LUT for a gradient that linearly transforms from c1 to c2.

Parameters:

• c1 (Color) – Input color as a string or 4-element RGBA list.

• c2 (Color) – Input color as a string or 4-element RGBA list.

Returns:

lut – Mayavi LUT for a gradient that linearly transforms from c1 to c2.

Return type:

npt.NDArray[np.uint8]

FullScreen()

This function makes the plot fullscreen.

IntRange

alias of Annotated[int]

PartScreen(width: float, height: float)

This function makes the plot width x height inches.

Parameters:

• width (float) – Width of the plot in inches.

• height (float) – Height of the plot in inches.

static SetColormapGradient(obj: Any, c1: str | tuple[float, float, float, float] | ndarray[Any, dtype[float64]], c2: str | tuple[float, float, float, float] | ndarray[Any, dtype[float64]])

Applies a linear, gradient colormap to the object. The colors in the gradient are c1 and c2.

Parameters:

• obj (Any) – mlab object to apply the colormap gradient to.

• c1 (Color) – Input color as a string or 4-element RGBA list.

• c2 (Color) – Input color as a string or 4-element RGBA list.

animate(animFunc: Callable[[], Generator[None, None, None]], outDir: str = 'MyMovie', fileName: str = 'images', save: bool = True, delay: int = 10) → mayavi.tools.animator.Animator | None

Parameters:

• animFunc (Callable[[], Generator[None, None, None]]) – Function that modifies what is displayed on the plot

• outDir (Path, optional) – Directory to save frames in: only used if save = True. (Default value = “MyMovie”)

• fileName (str, optional) – Filename for the frames: only used if save = True. (Default value = “images”)

• save (bool, optional) – Controls whether the function saves the frames of the animation or not. (Default value = True)

• delay (IntRange in [10,100000], optional) – Controls the delay used by mlab.animate. (Default value = 10)

Returns:

animator – Returns an animator object only if save = False. Otherwise, there is no return value.

Return type:

Optional[mayavi.tools.animator.Animator]

contour3d(*args: Any, **kwargs: Any) → mayavi.modules.iso_surface.IsoSurface

Call mlab’s contour3d on the class’s figure.

Parameters:

• *args (Any) – args passed on to contour3d

• **kwargs (Any) – kwargs passed on to mesh after being processed by _ProcessKwargs.

• Returns –

• -------- –

• contour (mayavi.modules.iso_surface.IsoSurface) – Coutour.

mesh(*args: Any, **kwargs: Any) → mayavi.modules.surface.Surface

Call mlab’s mesh on the class’s figure.

Parameters:

• *args (Any) – args passed on to mesh

• **kwargs (Any) – kwargs passed on to mesh after being processed by _ProcessKwargs.

• Returns –

• -------- –

• mesh (mayavi.modules.surface.Surface) – Mesh as a mayavi surface.

plot3d(*args: Any, **kwargs: Any) → mayavi.modules.surface.Surface

Call mlab’s plot3d on the class’s figure.

Parameters:

• *args (Any) – args passed on to plot3d

• **kwargs (Any) – kwargs passed on to plot3d after being processed by _ProcessKwargs.

• Returns –

• -------- –

• surf (mayavi.modules.surface.Surface) – Line as a Mayavi surface.

points3d(*args: Any, **kwargs: Any) → mayavi.modules.glyph.Glyph

Call mlab’s points3d on the class’s figure.

Parameters:

• *args (Any) – args passed on to points3d

• **kwargs (Any) – kwargs passed on to points3d after being processed by _ProcessKwargs.

• Returns –

• -------- –

• points3d (mayavi.modules.glyph.Glyph) – Glyphs corresponding to the points.

quiver3d(*args: Any, **kwargs: Any) → mayavi.modules.vectors.Vectors

Call mlab’s quiver3d on the class’s figure.

Parameters:

• *args (Any) – args passed on to quiver3d

• **kwargs (Any) – kwargs passed on to quiver3d after being processed by _ProcessKwargs.

• Returns –

• -------- –

• vectors (mayavi.modules.vectors.Vectors) – Vectors associated with the quiver3d plot.

save(fileName: str, fileType: Literal['png', 'jpg', 'bmp', 'tiff', 'ps', 'eps', 'pdf', 'rib', 'oogl', 'iv', 'wrl', 'vrml', 'obj', 'x3d', 'pov', None] | None = None)

This function saves the figure.

Parameters:

• fileName (Path) – Filename where the figure should be saved.

• fileType (Literal["png", "jpg", "bmp", "tiff", "ps", "eps", "pdf", "rib", "oogl", "iv", "wrl", "vrml", "obj", "x3d", "pov", None], optional) – File suffix to use. If None, then the suffix will be inferred from the file name. (Default value = None)

show()

Re-draw the class’s figure.

property show_axes

Adds axes indicator.

surf(*args: Any, **kwargs: Any) → mayavi.modules.surface.Surface

Call mlab’s surf on the class’s figure.

Parameters:

• *args (Any) – args passed on to surf

• **kwargs (Any) – kwargs passed on to surf after being processed by _ProcessKwargs.

• Returns –

• -------- –

• surf (mayavi.modules.surface.Surface) – Surface.

view(*args: Any, **kwargs: Any)

Call mlab’s view on the class’s figure.

Parameters:

• *args (Any) – args passed on to view

• **kwargs (Any) – kwargs passed on to view

tfc.utils.PlotlyMakePlot module

class tfc.utils.PlotlyMakePlot.MakePlot(xlabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str], ylabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str], titles: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str] | None = None, zlabs: _SupportsArray[dtype[str_]] | _NestedSequence[_SupportsArray[dtype[str_]]] | str | _NestedSequence[str] | None = None)

Bases: object

A MakePlot class for Plotly. This class wraps common Plotly functions to ease figure creation.

Box(row: int | None = None, col: int | None = None, **kwargs) → Box

Creates a plotly box on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graph_objects.Box

Returns:

box

Return type:

plotly.graphic_objects.Box

Contour(row: int | None = None, col: int | None = None, **kwargs) → Contour

Creates a plotly contour on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graphic_objects.Contour

Returns:

contour

Return type:

plotly.graph_objects.Contour

FullScreen()

Make the plot full screen.

Histogram(row: int | None = None, col: int | None = None, **kwargs) → Histogram

Creates a plotly histogram on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graph_objects.Histogram

Returns:

hist

Return type:

plotly.graphic_objects.Histogram

NormalizeColorScale(types: list[str] = [], data: str | None = None, cmax: float | None = None, cmin: float | None = None)

Normalizes the color scale for the plots whose type is in the types list.

If cmax and/or cmin are given, then the data variable is not used, and all plots whose type is in the types list are assigned that cmax/cmin value.

If cmax and cmin are not specified, then they are set by taking the max and min of the data that matches the data variable in all plots whose type is in the typpes list.

Parameters:

• types (List[str]) – Plot types to set cmax and cmin for.

• data (Optional[str]) – Data type to use to calculate cmax and cmin if not already specified. (Default value = None)

• cmax (Optional[float]) – cmax value to use when setting the colorscale

• cmin (Optional[float]) – cmin value to use when setting the colorscale

PartScreen(width: float, height: float, units: Literal['in', 'mm', 'px'] = 'in')

Make the plot size equal to width x height.

Parameters:

• width (float) – Width of the plot

• height (float) – Height of the plot

• units (Literal["in","mm","px"], optional) – Units width and height are given in. (Default value = inches)

Scatter(row: int | None = None, col: int | None = None, **kwargs) → Scatter

Creates a plotly scatter on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graph_objects.Scatter

Returns:

scatter

Return type:

plotly.graphic_objects.Scatter

Scatter3d(row: int | None = None, col: int | None = None, **kwargs) → Scatter3d

Creates a 3d plotly scatter on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graphic_objects.Scatter3d

Returns:

scatter

Return type:

plotly.graph_objects.Scatter3d

Surface(row: int | None = None, col: int | None = None, **kwargs) → Surface

Creates a plotly surface on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – Keyword arguments passed on to plotly.graphic_objects.Surface

Returns:

surf

Return type:

plotly.graph_objects.Surface

UploadToPlotly(username: str, apiKey: str, fileName: str, autoOpen: bool = False)

Upload your plot to Plotly.

Parameters:

• username (str) – Plotly username

• apiKey (str) – Plotly api_key

• fileName (str) – Name of the file to save the plot as.

• autoOpen (bool, optional) – If true, plot will open in browser after saving. (Default value = False)

Violin(row: int | None = None, col: int | None = None, **kwargs) → Violin

Creates a plotly violin on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graph_objects.Violin

Returns:

violin

Return type:

plotly.graphic_objects.Violin

Volume(row: int | None = None, col: int | None = None, **kwargs) → Volume

Creates a plotly volume on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• **kwargs (dict, optional) – keyword arguments passed on to plotly.graph_objects.Volume

Returns:

volume

Return type:

plotly.graphic_objects.Volume

save(fileName: str, tight: bool = True, fileType: Literal['pdf', 'jpg', 'png', 'svg', 'eps', 'html', None] | None = None, **kwargs)

Saves the figure using the type specified. If HTML is specified, the figure will be saved as a dynamic html file. All other file types are static.

Parameters:

• fileName (Path) – File name to save the figure as.

• tight (boolean, optional) – If the fileType is pdf or png and this value is true, then a tool is used to eliminate whitespace. pdfCropMargins is used for PDFs and PIL is used for png’s. (Default value = True)

• fileType (Literal["pdf","jpg","png","svg","eps","html",None], optional) – File suffix to use. If None, then the suffix will be inferred from the suffix of the fileName. (Default value = None)

• **kwargs (dict, optional) – Keyword arguments passed onto fig.write_image or fig.write_html, depending on fileType.

show(**kwargs)

Calls the figure’s show method.

Parameters:

**kwargs (dict, optional) – keyword arguments passed on to fig.show

view(azimuth: float, elevation: float, row: int | None = None, col: int | None = None, viewDistance: float | None = None)

Change the view on the subplot specified by row and col or on the main figure if not using subplots.

Parameters:

• azimuth (float) – Azimuth value in degrees

• elevation (float) – Elevation value in degrees

• row (Optional[uint]) – Subplot row (Default value = None)

• col (Optional[uint]) – Subplot column (Default value = None)

• viewDistance (Optional[float]) – Distance from camera to plot. (Default value = self.viewDistance)

tfc.utils.TFCUtils module

class tfc.utils.TFCUtils.ComponentConstraintDict

Bases: TypedDict

name: str

node0: str

node1: str

class tfc.utils.TFCUtils.ComponentConstraintGraph(N: list[str], E: list[ComponentConstraintDict])

Bases: object

Creates a graph of all valid ways in which component constraints can be embedded.

SaveGraphs(outputDir: str, allGraphs: bool = False, savePDFs: bool = False) → None

Saves the graphs. The graphs are saved in a clickable HTML structure. They can also be saved as PDFs.

Parameters:

• outputDir (Path) – Output directory to save in.

• allGraphs (bool, optional) – Boolean that conrols whether all graphs are saved or just valid graphs. (Default value = False)

• savePDFs (bool, optional) – Boolean that controls whether the graphs are also saved as PDFs. (Default value = False)

tfc.utils.TFCUtils.LS(zXi: PyTree, res: Callable, *args: Any, constant_arg_nums: list[int] = [], J: Callable[[...], Array] | None = None, method: Literal['pinv', 'lstsq'] = 'pinv', timer: bool = False, timerType: str = 'process_time', holomorphic: bool = False) → PyTree | tuple[PyTree, float]

JITed least squares. This function takes in an initial guess of zeros, zXi, and a residual function, res, and linear least squares to minimize the res function using the parameters xi.

Parameters:

• zXi (PyTree) – Unknown parameters to be found using least-squares.

• res (Callable) – Residual function (also known as the loss function) with signature res(xi: PyTree, *args:Any, **kwargs:Any). Note, the first argument does not need to be named xi, this is just illustrative.

• *args (Any) – Any additional arguments taken by res other than the first PyTree argument.

• constant_arg_nums (List[int], optional) – These arguments will be removed from the residual function and treated as constant. See pejit for more details.

• J (Optional[Callable[...,np.ndarray]]) – User specified Jacobian function. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)

• method (Literal["pinv","lstsq"], optional) – Method for least-squares inversion. (Default value = “pinv”) * pinv - Use np.linalg.pinv * lstsq - Use np.linalg.lstsq

• timer (bool, optional) – Boolean that chooses whether to time the code or not. (Default value = False). Note that setting to true adds a slight increase in runtime. As one iteration of the non-linear least squares is run first to avoid timining the JAX trace.

• timerType (str, optional) – Any timer from the time module. (Default value = “process_time”)

• holomorphic (bool, optional) – Indicates whether residual function is promised to be holomorphic. (Default value = False)

Returns:

• xi (pytree or array-like) – Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.

• time (float) – Computation time as calculated by timerType specified. This output is only returned if timer = True.

class tfc.utils.TFCUtils.LsClass(zXi: PyTree, res: Callable, *args: Any, constant_arg_nums: list[int] = [], J: Callable[[...], Array] | None = None, method: Literal['pinv', 'lstsq'] = 'pinv', timer: bool = False, timerType: str = 'process_time', holomorphic: bool = False)

Bases: object

JITed linear least-squares class. Like the LS function, but it is in class form so that the run methd can be called multiple times without re-JITing. See LS for more details.

run(zXi: PyTree, *args: Any) → PyTree | tuple[PyTree, float]

Runs the JIT-ed least-squares function and times it if desired.

Parameters:

• zXi (PyTree) – Unknown parameters to be found using least-squares.

• *args (Any) – Any additional arguments taken by res other than xi.

Returns:

• xi (PyTree) – Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.

• time (float, optional) – Computation time as calculated by timerType specified. This output is only returned if timer = True.

tfc.utils.TFCUtils.NLLS(xiInit: PyTree, res: Callable, *args: Any, constant_arg_nums: list[int] = [], J: Callable[[...], Array] | None = None, cond: Callable[[PyTree], bool] | None = None, body: Callable[[PyTree], PyTree] | None = None, tol: float = 1e-13, maxIter: int = 50, method: Literal['pinv', 'lstsq'] = 'pinv', timer: bool = False, printOut: bool = False, printOutEnd: str = '\n', timerType: str = 'process_time', holomorphic: bool = False) → tuple[PyTree, int] | tuple[PyTree, int, float]

JIT-ed non-linear least squares. This function takes in an initial guess, xiInit (initial values of xi), and a residual function, res, and performs a nonlinear least squares to minimize the res function using the parameters xi. The conditions on terminating the nonlinear least-squares are: 1. max(abs(res)) < tol 2. max(abs(dxi)) < tol, where dxi is the change in xi from the last iteration. 3. Number of iterations > maxIter.

Parameters:

• xiInit (pytree or array-like) – Initial guess for the unkown parameters.

• res (function) – Residual function (also known as the loss function) with signature res(xi,*args).

• *args (iterable) – Any additional arguments taken by res other than xi.

• constant_arg_nums (List[int], optional) – These arguments will be removed from the residual function and treated as constant. See pejit for more details.

• J (function) – User specified Jacobian. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)

• cond (Optional[Callable[[PyTree],bool]]) – User specified condition function. If None, then the default cond function is used which checks the three termination criteria provided in the class description. (Default value = None)

• body (Optional[Callable[[PyTree],PyTree]]) – User specified while-loop body function. If None, then use the default body function which updates xi using a NLLS interation and the method provided. (Default value = None)

• tol (float) – Tolerance used in the default termination criteria: see class description for more details. (Default value = 1e-13)

• maxIter (int, optional) – Maximum number of iterations. (Default value = 50)

• method (Literal["pinv","lstsq"], optional) – Method for least-squares inversion. (Default value = “pinv”) * pinv - Use np.linalg.pinv * lstsq - Use np.linalg.lstsq

• timer (bool, optional) – Boolean that chooses whether to time the code or not. (Default value = False). Note that setting to true adds a slight increase in runtime. As one iteration of the non-linear least squares is run first to avoid timining the JAX trace.

• printOut (bool, optional) – Controls whether the NLLS prints out information each interaton or not. The printout consists of the iteration and max(abs(res)) at each iteration. (Default value = False)

• printOutEnd (str, optional) – Value of keyword argument end passed to the print statement used in printOut. (Default value = “\n”)

• timerType (str, optional) – Any timer from the time module. (Default value = “process_time”)

• holomorphic (bool, optional) – Indicates whether residual function is promised to be holomorphic. (Default value = False)

Returns:

• xi (PyTree) – Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.

• it (int) – Number of NLLS iterations performed.

• time (float) – Computation time as calculated by timerType specified. This output is only returned if timer = True.

class tfc.utils.TFCUtils.NllsClass(xiInit: PyTree, res: Callable, *args: Any, constant_arg_nums: list[int] = [], J: Callable[[...], Array] | None = None, cond: Callable[[PyTree], bool] | None = None, body: Callable[[PyTree], PyTree] | None = None, tol: float = 1e-13, maxIter: int = 50, method: Literal['pinv', 'lstsq'] = 'pinv', timer: bool = False, printOut: bool = False, printOutEnd: str = '\n', timerType: str = 'process_time', holomorphic: bool = False)

Bases: object

JITed nonlinear least squares class. Like the NLLS function, but it is in class form so that the run methd can be called multiple times without re-JITing

run(xiInit: PyTree, *args: Any) → tuple[PyTree, int] | tuple[PyTree, int, float]

Runs the JIT-ed nonlinear least-squares function and times it if desired.

Parameters:

• xiInit (PyTree) – Initial guess for the unkown parameters.

• *args (Any) – Any additional arguments taken by res other than xi.

Returns:

• xi (PyTree) – Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.

• it (int) – Number of NLLS iterations performed..

• time (float) – Computation time as calculated by timerType specified. This output is only returned if timer = True.

tfc.utils.TFCUtils.ScaledQrLs(A: Array, B: Array) → tuple[Array, Array]

This function performs least-squares using a scaled QR method.

Parameters:

• A (np.ndarray) – A matrix in A*x = B.

• B (np.ndarray) – B matrix in A*x = B.

Returns:

• x (np.ndarray) – Solution to A*x = B solved using a scaled QR method.

• cn (np.ndarray) – Condition number.

class tfc.utils.TFCUtils.TFCDict(*args: Any)

Bases: OrderedDict

This is the TFC dictionary class. It extends an OrderedDict and adds a few methods that enable:

• Adding dictionaries with the same keys together

• Turning a dictionary into a 1-D array

• Turning a 1-D array into a dictionary

block_until_ready() → TFCDict

Mimics block_until_ready for jax arrays. Used to halt the program until the computation that created the dictionary is finished.

Returns:

This TFCDict.

Return type:

TFCDict

getSlices() → None

Function that creates slices for each of the keys in the dictionary.

toArray() → Array

Send dictionary to a flat JAX array.

Returns:

This dictionary as a flat JAX array.

Return type:

np.ndarray

toDict(arr: Array) → TFCDict

Send a flat JAX array to a TFCDict with the same keys.

Parameters:

arr (ndarray) – Flat JAX array to convert to TFCDict. Must have the same number of elements as total number of elements in the dictionary.

Returns:

JAX array as a TFCDict

Return type:

TFCDict

update(*args: Any) → None

Overload the update method to update the _keys variable as well.

Parameters:

*args (Any) – Same as *args for the update method of ordered dict.

class tfc.utils.TFCUtils.TFCDictRobust(*args: Any)

Bases: OrderedDict

This class is like the TFCDict class, but it handles non-flat arrays.

block_until_ready() → TFCDictRobust

Mimics block_until_ready for jax arrays. Used to halt the program until the computation that created the dictionary is finished.

Returns:

This TFCDictRobust

Return type:

TFCDictRobust

getSlices() → None

Function that creates slices for each of the keys in the dictionary.

toArray() → Array

Send dictionary to a flat JAX array.

Returns:

This dictionary as a flat JAX array.

Return type:

np.ndarray

toDict(arr: Array) → TFCDictRobust

Send a flat JAX array to a TFCDictRobust with the same keys.

Parameters:

arr (np.ndarray) – Flat JAX array to convert to TFCDictRobust. Must have the same number of elements as total number of elements in the dictionary.

Returns:

JAX array as a TFCDictRobust

Return type:

TFCDictRobust

update(*args: Any) → None

Overload the update method to update the _keys variable as well.

Parameters:

*args (Any) – Same as *args for the update method on ordered dict.

class tfc.utils.TFCUtils.TFCPrint

Bases: object

This class is used to print to the terminal in color.

static Error(stringIn: str)

This function prints errors. It prints the text in ‘stringIn’ in bright red and exits the program.

Parameters:

stringIn (str) – error string

static Warning(stringIn: str)

This function prints warnings. It prints the text in ‘stringIn’ in bright yellow.

Parameters:

stringIn (str) – warning string

tfc.utils.TFCUtils.criuCheckpoint(dir: str = 'criu_checkpoint', user_mode: bool = False)

Use CRIU to create a checkpoint of your program. WARNING: You must ensure no external sockets are used, i.e., via matplotlib. WARNING: GPU memory cannot be mapped for all GPUs. WARNING: user_mode is a work in progress and is not full featured yet.

Parameters:

• dir (str) – Directory where the CRIU checkpoint will be created.

• user_mode (bool) – Whether to use user mode or not. (Default value = False)

tfc.utils.TFCUtils.egrad(g: Callable[[...], Any], j: int = 0)

This function mimics egrad from autograd.

Parameters:

• g (Callable[..., Any]) – Function to take the derivative of.

• j (uint, optional) – Parameter with which to take the derivative with respect to. (Default value = 0)

Returns:

wrapped – Derivative function

Return type:

function

tfc.utils.TFCUtils.egradRobust(g: Callable[[...], Any], j: int = 0)

This function mimics egrad from autograd, but can also handle dictionaries.

Parameters:

• g (function) – Function to take the derivative of.

• j (integer, optional) – Parameter with which to take the derivative with respect to. (Default value = 0)

Returns:

wrapped – Derivative function

Return type:

function

tfc.utils.TFCUtils.nlls_id_print(it: int, x, end: str = '\n')

tfc.utils.TFCUtils.pe(*args: Any, constant_arg_nums: list[int] = []) → Any

Decorator that returns a function evaluated such that the arg numbers specified in constant_arg_nums and all functions that utilizes only those arguments are treated as compile time constants.

Parameters:

• *args (Any) – Arguments for the function that pe is applied to.

• constant_arg_nums (List[int], optional) – The arguments whose values and functions that depend only on these values should be treated as cached constants.

Returns:

• f (Any) – The new function whose constant_arg_num arguments have been removed. The jaxpr of this function has the constant_arg_num values and all functions that depend on those values cached as constants.

• Usage

• —–

• @pe(*args, constant_arg_nums=[0])

• def f(x,xi) – # Function stuff here

• # Returns an f(xi) with x treated as constant

tfc.utils.TFCUtils.pejit(*args: Any, constant_arg_nums: list[int] = [], **kwargs) → Any

Works like pe, but also JITs the returned function. See pe for more details.

Parameters:

*args: Any

Arguments for the function that pe is applied to.

constant_arg_nums: List[int], optional

The arguments whose values and functions that depend only on these values should be treated as cached constants.

**kwargs: Any

Keyword arguments passed on to JIT.

Returns:

f: Any

The new function whose constant_arg_num arguments have been removed. The jaxpr of this function has the constant_arg_num values and all functions that depend on those values cached, and they are treated as compile time constants.

tfc.utils.TFCUtils.step(x: Array) → Array

This is the unit step function, but the deriative is defined and equal to 0 at every point.

Parameters:

x (np.ndarray) – Array to apply step to.

Returns:

step_x – step(x)

Return type:

np.ndarray

Module contents