1"""A variety of helper functions.
  2
  3Notably, includes functions for generalized reading and writing from
  4mesh files. This is necessary due to the variety of interfaces that
  5may be easiest to work with, given a task.
  6
  7A `meshio.Mesh` object is useful for extracting point/cell data
  8according to the event horizon mesh output from `gel.scripts.get_veh`.
  9`nodal_values_to_meshio` achieves this task.
 10
 11"Full-shape" .xdmf files are created from write_checkpoint methods, and
 12are suitable for reading and writing directly in and out of FEniCS
 13without additional machinery necessary. Associated functions include
 14`save_shape` and `load_shape`.
 15
 16"Nodal" .xdmf files are best for post-processing with Paraview as they
 17have a compatible embedding of connectivity information and are smaller
 18files. However, some floating-point precision loss has been observed.
 19Associated functions for working with function DoFs in this manner
 20include `save_fenics_fcn_nodally` and `nodal_values_to_fenics_fcn`.
 21
 22# API
 23"""
 24from .header import *
 25import os
 26import argparse
 27from functools import cmp_to_key
 28
 29
 30ghrp = lambda data_directory : (
 31    lambda fname : os.path.join(data_directory, fname)
 32)
 33"""Given directory `data_directory`, return mapping filename to path."""
 34
 35
 36def get_common_parser(*args, **kwargs):
 37    """Returns `argparse.ArgumentParser` equipped with common args.
 38
 39    * `args`: positional arguments to `argparse.ArgumentParser`
 40    * `kwargs`: keyword arguments to `argparse.ArgumentParser`
 41
 42    Included arguments:
 43    ```
 44    -c CELL_DATA_DIR
 45    -f FORMULATION
 46    -k MODULUS_RATIO
 47    -b BOUNDS
 48    -l LOAD_STEPS
 49    -p PRECONDITIONER
 50    --bci CELL_SURF_MESH
 51    --bco OUTER_SURF_MESH
 52    ```
 53    """
 54    if "formatter_class" not in kwargs:
 55        kwargs["formatter_class"] = argparse.ArgumentDefaultsHelpFormatter
 56
 57    parser = argparse.ArgumentParser(*args, **kwargs)
 58
 59    parser.add_argument(
 60        "-c",
 61        metavar="CELL_DATA_DIR",
 62        default="cell_data",
 63        help="directory containing gel geometry"
 64    )
 65    parser.add_argument(
 66        "-f",
 67        metavar="FORMULATION",
 68        default="beta",
 69        help="form of strain energy density"
 70    )
 71    parser.add_argument(
 72        "-k",
 73        type=float,
 74        metavar="MODULUS_RATIO",
 75        default=1,
 76        help="D1/C1 ratio"
 77    )
 78    parser.add_argument(
 79        "-b",
 80        nargs=2,
 81        type=float,
 82        metavar="BOUNDS",
 83        default=[-2.0, 20.0],
 84        help="bounds for inverse model, or built-in for some formulations"
 85    )
 86    parser.add_argument(
 87        "-l",
 88        type=int,
 89        metavar="LOAD_STEPS",
 90        default=1,
 91        help="forward model BC load step count"
 92    )
 93    parser.add_argument(
 94        "-p",
 95        type=str,
 96        metavar="PRECONDITIONER",
 97        default="hypre_amg",
 98        help="preconditioner for forward model Newton-Raphson linear solves"
 99    )
100    parser.add_argument(
101        "--bci",
102        type=str,
103        metavar="CELL_SURF_MESH",
104        default=None,
105        help="filename of meshio-compatible mesh with cell surface nodes and u"
106    )
107    parser.add_argument(
108        "--bco",
109        type=str,
110        metavar="OUTER_SURF_MESH",
111        default=None,
112        help="filename of meshio-compatible mesh with outer nodes and u"
113    )
114
115    return parser
116
117
118@cmp_to_key
119def lexico_sort(coord0, coord1):
120    for i in range(len(coord0)):
121        cmp = coord0[i] - coord1[i]
122        if cmp != 0:
123            return cmp
124    return 0
125
126
127def save_fenics_fcn_nodally(fcn, filename, fcn_name, long_fcn_name=None):
128    """Writes a nodal .xdmf file with given function.
129
130    * `fcn`: FEniCS FE function to write
131    * `filename`: str filename ending in ".xdmf"
132    * `fcn_name`: str short name of the function
133    * `long_fcn_name`: str long name of the function (default same as
134    `fcn_name`)
135
136    Side-effects: writes files `filename` with .xdmf and .h5 endings
137    """
138    output_file = XDMFFile(filename)
139    output_file.parameters["flush_output"] = True
140    output_file.parameters["functions_share_mesh"] = True
141
142    if long_fcn_name is None:
143        long_fcn_name = fcn_name
144
145    fcn.rename(fcn_name, long_fcn_name)
146
147    output_file.write(fcn, 0)
148
149
150def nodal_values_to_meshio(filename):
151    """Reads an .xdmf in "nodal" form to return a `meshio.Mesh` with
152    all point and cell data.
153
154    Note that the `meshio.xdmf.TimeSeriesReader` is involved, but only
155    time 0 is read in accordance with this libraries output style.
156
157    * `filename`: str filename ending in ".xdmf", must be in "nodal"
158    form
159    """
160    # Read in nodal DOFs
161    with meshio.xdmf.TimeSeriesReader(filename) as reader:
162        points, cells = reader.read_points_cells()
163        _, point_data, cell_data = reader.read_data(0)
164
165    return meshio.Mesh(
166        points,
167        cells,
168        point_data=point_data,
169        cell_data=cell_data
170    )
171
172
173def nodal_values_to_fenics_fcn(
174        geo,
175        nodal_values_filename,
176        field_name="mod_repr"
177    ):
178    """Reads a nodal .xdmf file to create a new **scalar** function.
179
180    * `geo`: `gel.geometry.Geometry` with matching hydrogel mesh
181    * `nodal_values_filename`: str filename ending in ".xdmf", must be
182    a "full_shape" form
183    * `field_name`: str the name of the scalar field in the file
184
185    Returns: FEniCS FE scalar function in 1st order Lagrange space
186    """
187    # Read in nodal DOFs
188    mesh = nodal_values_to_meshio(nodal_values_filename)
189    points, cells, point_data, cell_data = \
190        mesh.points, mesh.cells, mesh.point_data, mesh.cell_data
191
192    # Corresponding function space
193    coordinates = geo.V0.tabulate_dof_coordinates()
194
195    # Helper functions
196    sorted_to_meshio = sorted(
197        range(len(points)), key = lambda i : lexico_sort(points[i])
198    )
199    sorted_to_fenics = sorted(
200        range(len(coordinates)), key = lambda i : lexico_sort(coordinates[i])
201    )
202
203    # Into FEniCS order
204    dofs_fenics_order = np.zeros(len(coordinates))
205    dofs_fenics_order[sorted_to_fenics] = \
206            point_data[field_name].flatten()[sorted_to_meshio]
207
208    # Make function
209    fcn = Function(geo.V0)
210    fcn.vector().set_local(dofs_fenics_order)
211    fcn.vector().apply("")
212    fcn.set_allow_extrapolation(True)
213
214    return fcn
215
216
217def save_shape(fcn, filename, fcn_name):
218    """Writes full-shape .xdmf file with given function.
219
220    * `fcn`: FEniCS FE function to write
221    * `filename`: str filename ending in ".xdmf"
222    * `fcn_name`: str short name of the function
223
224    Side-effects: writes files `filename` with .xdmf and .h5 endings
225    """
226    shape_file = XDMFFile(filename)
227    shape_file.parameters["flush_output"] = True
228    shape_file.parameters["functions_share_mesh"] = True
229
230    shape_file.write_checkpoint(
231        fcn,
232        fcn_name,
233        0,
234        XDMFFile.Encoding.HDF5,
235        False
236    )
237
238
239def load_shape(fcn_space, filename, fcn_name):
240    """Reads a full-shape .xdmf file to create a new FEniCS FE function.
241
242    * `fcn_space`: FEniCS FunctionSpace that the function was originally
243    created in
244    * `filename`: str filename ending in ".xdmf", must be in
245    "full_shape" form
246    * `fcn_name`: str the name of the function in the file
247
248    Returns: FEniCS FE function
249    """
250    fcn = Function(fcn_space)
251    shape_file = XDMFFile(filename)
252    shape_file.read_checkpoint(fcn, fcn_name, 0)
253
254    return fcn