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    --u-init U_INIT
 53    ```
 54    """
 55    if "formatter_class" not in kwargs:
 56        kwargs["formatter_class"] = argparse.ArgumentDefaultsHelpFormatter
 57
 58    parser = argparse.ArgumentParser(*args, **kwargs)
 59
 60    parser.add_argument(
 61        "-c",
 62        metavar="CELL_DATA_DIR",
 63        default="cell_data",
 64        help="directory containing gel geometry"
 65    )
 66    parser.add_argument(
 67        "-f",
 68        metavar="FORMULATION",
 69        default="beta",
 70        help="form of strain energy density"
 71    )
 72    parser.add_argument(
 73        "-k",
 74        type=float,
 75        metavar="MODULUS_RATIO",
 76        default=1,
 77        help="D1/C1 ratio"
 78    )
 79    parser.add_argument(
 80        "-b",
 81        nargs=2,
 82        type=float,
 83        metavar="BOUNDS",
 84        default=[-2.0, 20.0],
 85        help="bounds for inverse model, or built-in for some formulations"
 86    )
 87    parser.add_argument(
 88        "-l",
 89        type=int,
 90        metavar="LOAD_STEPS",
 91        default=1,
 92        help="forward model BC load step count"
 93    )
 94    parser.add_argument(
 95        "-p",
 96        type=str,
 97        metavar="PRECONDITIONER",
 98        default="hypre_amg",
 99        help="preconditioner for forward model Newton-Raphson linear solves"
100    )
101    parser.add_argument(
102        "--bci",
103        type=str,
104        metavar="CELL_SURF_MESH",
105        default=None,
106        help="filename of meshio-compatible mesh with cell surface nodes and u"
107    )
108    parser.add_argument(
109        "--bco",
110        type=str,
111        metavar="OUTER_SURF_MESH",
112        default=None,
113        help="filename of meshio-compatible mesh with outer nodes and u"
114    )
115    parser.add_argument(
116        "--u-init",
117        type=str,
118        metavar="U_INIT",
119        default=None,
120        help="filename of full-shape .xdmf file with initial displacements"
121    )
122
123    return parser
124
125
126@cmp_to_key
127def lexico_sort(coord0, coord1):
128    for i in range(len(coord0)):
129        cmp = coord0[i] - coord1[i]
130        if cmp != 0:
131            return cmp
132    return 0
133
134
135def save_fenics_fcn_nodally(fcn, filename, fcn_name, long_fcn_name=None):
136    """Writes a nodal .xdmf file with given function.
137
138    * `fcn`: FEniCS FE function to write
139    * `filename`: str filename ending in ".xdmf"
140    * `fcn_name`: str short name of the function
141    * `long_fcn_name`: str long name of the function (default same as
142    `fcn_name`)
143
144    Side-effects: writes files `filename` with .xdmf and .h5 endings
145    """
146    output_file = XDMFFile(filename)
147    output_file.parameters["flush_output"] = True
148    output_file.parameters["functions_share_mesh"] = True
149
150    if long_fcn_name is None:
151        long_fcn_name = fcn_name
152
153    fcn.rename(fcn_name, long_fcn_name)
154
155    output_file.write(fcn, 0)
156
157
158def nodal_values_to_meshio(filename):
159    """Reads an .xdmf in "nodal" form to return a `meshio.Mesh` with
160    all point and cell data.
161
162    Note that the `meshio.xdmf.TimeSeriesReader` is involved, but only
163    time 0 is read in accordance with this libraries output style.
164
165    * `filename`: str filename ending in ".xdmf", must be in "nodal"
166    form
167    """
168    # Read in nodal DOFs
169    with meshio.xdmf.TimeSeriesReader(filename) as reader:
170        points, cells = reader.read_points_cells()
171        _, point_data, cell_data = reader.read_data(0)
172
173    return meshio.Mesh(
174        points,
175        cells,
176        point_data=point_data,
177        cell_data=cell_data
178    )
179
180
181def nodal_values_to_fenics_fcn(
182        geo,
183        nodal_values_filename,
184        field_name="mod_repr"
185    ):
186    """Reads a nodal .xdmf file to create a new **scalar** function.
187
188    * `geo`: `gel.geometry.Geometry` with matching hydrogel mesh
189    * `nodal_values_filename`: str filename ending in ".xdmf", must be
190    a "full_shape" form
191    * `field_name`: str the name of the scalar field in the file
192
193    Returns: FEniCS FE scalar function in 1st order Lagrange space
194    """
195    # Read in nodal DOFs
196    mesh = nodal_values_to_meshio(nodal_values_filename)
197    points, cells, point_data, cell_data = \
198        mesh.points, mesh.cells, mesh.point_data, mesh.cell_data
199
200    # Corresponding function space
201    coordinates = geo.V0.tabulate_dof_coordinates()
202
203    # Helper functions
204    sorted_to_meshio = sorted(
205        range(len(points)), key = lambda i : lexico_sort(points[i])
206    )
207    sorted_to_fenics = sorted(
208        range(len(coordinates)), key = lambda i : lexico_sort(coordinates[i])
209    )
210
211    # Into FEniCS order
212    dofs_fenics_order = np.zeros(len(coordinates))
213    dofs_fenics_order[sorted_to_fenics] = \
214            point_data[field_name].flatten()[sorted_to_meshio]
215
216    # Make function
217    fcn = Function(geo.V0)
218    fcn.vector().set_local(dofs_fenics_order)
219    fcn.vector().apply("")
220    fcn.set_allow_extrapolation(True)
221
222    return fcn
223
224
225def save_shape(fcn, filename, fcn_name):
226    """Writes full-shape .xdmf file with given function.
227
228    * `fcn`: FEniCS FE function to write
229    * `filename`: str filename ending in ".xdmf"
230    * `fcn_name`: str short name of the function
231
232    Side-effects: writes files `filename` with .xdmf and .h5 endings
233    """
234    shape_file = XDMFFile(filename)
235    shape_file.parameters["flush_output"] = True
236    shape_file.parameters["functions_share_mesh"] = True
237
238    shape_file.write_checkpoint(
239        fcn,
240        fcn_name,
241        0,
242        XDMFFile.Encoding.HDF5,
243        False
244    )
245
246
247def load_shape(fcn_space, filename, fcn_name):
248    """Reads a full-shape .xdmf file to create a new FEniCS FE function.
249
250    * `fcn_space`: FEniCS FunctionSpace that the function was originally
251    created in
252    * `filename`: str filename ending in ".xdmf", must be in
253    "full_shape" form
254    * `fcn_name`: str the name of the function in the file
255
256    Returns: FEniCS FE function
257    """
258    fcn = Function(fcn_space)
259    shape_file = XDMFFile(filename)
260    shape_file.read_checkpoint(fcn, fcn_name, 0)
261
262    return fcn