Physics-informed graph neural network (GNN) emulation of soft-tissue mechanics
"""
File: constitutive_law.py
Author: David Dalton
Description: Various setups before training, e.g. constitutive law
"""

####################################################
## Package imports
####################################################

import jax.numpy as jnp
from jax import device_put
import os
import sys

####################################################
## Transformation functions for stabilising strain
## energy - see Section 2.4.3 of paper for details
####################################################

Jmin   = 0.001 # minimum value of J which we allow
Jtrans = 0.05  # point at which transformation function kicks in

Jdiff = Jtrans - Jmin

beta_1 = 1. / Jdiff
beta_2 = jnp.log(Jdiff) - Jtrans/Jdiff
beta_3 = Jmin

exp_fn              = lambda J: jnp.exp(beta_1*J + beta_2) + beta_3
J_transformation_fn = lambda J: jnp.where(J > Jtrans, J, exp_fn(J))

I1_trans    = 10 # point at which transformation function kicks in
tanh_fn     = lambda I1: jnp.tanh(I1 - I1_trans) + I1_trans
I1_trans_fn = lambda I1: jnp.where(I1 < I1_trans, I1, tanh_fn(I1))

I4f_trans    = 8 # point at which transformation function kicks in
tanh_fn      = lambda I4f: jnp.tanh(I4f - I4f_trans) + I4f_trans
I4f_trans_fn = lambda I4f: jnp.where(I4f < I4f_trans, I4f, tanh_fn(I4f))


####################################################
## Define constitutive law 
####################################################

# incompressibility penalty parameter
kappa = 25.

def HolzapfelOgden(params, F, J, f0):
    """Transversely-isotropic HO strain energy density function

    Inputs:
    -----------
    params: jnp.array
       Material parameter vector
    F     : jnp.array
       Deformation gradient
    J     : jnp.array
       Determinant of F
    f0:   : jnp.array
       Myofibre orientation vector

    Returns:
    ----------
    sed: float
       Strain energy density
    """

    # extract Lame material parameters
    a, b, af, bf, = params

    # right Cauchy-Green tensor (using F-bar method)
    C = jnp.power(J, -2./3.) * jnp.matmul(F.T, F)

    # first invariant and I4f pseudo-invariant
    I1 = jnp.trace(C)
    I4f = jnp.matmul(jnp.matmul(f0.T, C), f0)

    # transform to stop very large values
    I1 = I1_trans_fn(I1)
    I4f = I4f_trans_fn(I4f)
     
    sed = (a/b) * (jnp.exp(b * (I1 - 3.)) - 1.) + (af/bf)*(jnp.exp(bf*(I4f - 1.)**2) - 1) + kappa * (J**2 - 1. - 2.*jnp.log(J))

    return sed

constitutive_law = HolzapfelOgden 

####################################################
## Define function to enforce essential boundary conditions
####################################################

abs_path= os.path.join(sys.path[0])

interior_points = device_put(jnp.load(f'{abs_path}/interior-points.npy')*1.).reshape(-1,1)

boundary_adjust_fn = lambda U: U*interior_points


####################################################
## Define external forces (boundary and traction forces)
####################################################

final_pressure_loading = 1.06658
body_force             = None

####################################################
## Define material parameter boundaries
####################################################

params_lb = jnp.array([.1,  1.,  1.5,  1.])
params_ub = jnp.array([.26, 4.2, 5.18, 4.46])

####################################################
## Sample parameters on log scale or uniform scale
####################################################

log_sampling = True

####################################################
## Set number of parameters to sample at each epoch
####################################################

epoch_size = 200



data/LeftVentricle/geometryData/constitutive_law.py
dodaltuin/soft-tissue-pignn
"""
File: constitutive_law.py
Author: David Dalton
Description: Various setups before training, e.g. constitutive law
"""

####################################################
## Package imports
####################################################

import jax.numpy as jnp
from jax import device_put
import os
import sys

####################################################
## Transformation functions for stabilising strain 
## energy - see Section 2.4.3 of paper for details
####################################################

Jmin = 0.001  # minimum value of J which we allow
Jtrans = 0.05 # point at which transformation function kicks in

Jdiff = Jtrans - Jmin

beta_1 = 1. / Jdiff
beta_2 = jnp.log(Jdiff) - Jtrans/Jdiff
beta_3 = Jmin

exp_fn = lambda J: jnp.exp(beta_1*J + beta_2) + beta_3
J_transformation_fn = lambda J: jnp.where(J > Jtrans, J, exp_fn(J))

I1_trans = 10 # point at which transformation function kicks in

tanh_fn = lambda I1: jnp.tanh(I1 - I1_trans) + I1_trans
I1_trans_fn = lambda I1: jnp.where(I1 < I1_trans, I1, tanh_fn(I1))

####################################################
## Define constitutive law 
####################################################

def NeoHookean(params, F, J, fibres=None):
    """NeoHookean strain energy density function

    Inputs:
    -----------
    params: jnp.array
       Material parameter vector
    F     : jnp.array
       Deformation gradient
    J     : jnp.array
       Determinant of F
    fibres: None
       Not used here as we assume isotropic material

    Returns:
    ----------
    sed: float
       Strain energy density
    """

    # extract Lame material parameters
    lambda_, mu = params

    # right Cauchy-Green tensor
    C = jnp.matmul(F.T, F)

    # first invariant
    Ic = jnp.trace(C)

    # transform to stop very large values
    Ic = I1_trans_fn(Ic)

    lnJ = jnp.log(J)

    sed = (mu/2.)*(Ic - 3.) - mu*lnJ + (lambda_/2.)*(lnJ**2)

    return sed

constitutive_law = NeoHookean

####################################################
## Define function to enforce essential boundary conditions
####################################################

abs_path= os.path.join(sys.path[0])

interior_points = device_put(jnp.load(f'{abs_path}/interior-points.npy')*1.)

boundary_adjust_fn = lambda U: U*interior_points


####################################################
## Define external forces (boundary and traction forces)
####################################################

final_pressure_loading = None
g      = 980
body_force             = jnp.array([0.]*(2) + [-g])

####################################################
## Define material parameter boundaries
####################################################

params_lb = jnp.array([35000.]*2)
params_ub = jnp.array([100000.]*2)

####################################################
## Sample parameters on log scale or uniform scale 
####################################################

log_sampling = True

####################################################
## Set number of parameters to sample at each epoch 
####################################################

epoch_size = 200






data/Liver/geometryData/constitutive_law.py
dodaltuin/soft-tissue-pignn
"""
File: constitutive_law.py
Author: David Dalton
Description: Various setups before training, e.g.constitutive law
"""

####################################################
## Package imports
####################################################

import jax.numpy as jnp
from jax import device_put
import os
import sys

####################################################
## Transformation functions for stabilising strain 
## energy - see Section 2.4.3 of paper for details
####################################################

Jmin = 0.001  # minimum value of J which we allow
Jtrans = 0.05 # point at which transformation function kicks in

Jdiff = Jtrans - Jmin

beta_1 = 1. / Jdiff
beta_2 = jnp.log(Jdiff) - Jtrans/Jdiff
beta_3 = Jmin

exp_fn = lambda J: jnp.exp(beta_1*J + beta_2) + beta_3
J_transformation_fn = lambda J: jnp.where(J > Jtrans, J, exp_fn(J))

I1_trans = 10 # point at which transformation function kicks in

tanh_fn = lambda I1: jnp.tanh(I1 - I1_trans) + I1_trans
I1_trans_fn = lambda I1: jnp.where(I1 < I1_trans, I1, tanh_fn(I1))

####################################################
## Define constitutive law 
####################################################

def NeoHookean(params, F, J, fibres=None):
    """NeoHookean strain energy density function

    Inputs:
    -----------
    params: jnp.array
       Material parameter vector
    F     : jnp.array
       Deformation gradient
    J     : jnp.array
       Determinant of F
    fibres: None
       Not used here as we assume isotropic material

    Returns:
    ----------
    sed: float
       Strain energy density
    """

    # extract Lame material parameters
    mu, lambda_ = params

    # right Cauchy-Green tensor
    C = jnp.matmul(F.T, F)

    # first invariant
    Ic = jnp.trace(C)

    # transform to stop very large values
    Ic = I1_trans_fn(Ic)

    lnJ = jnp.log(J)

    sed = (mu/2.)*(Ic - 3.) - mu*lnJ + (lambda_/2.)*(lnJ**2)

    return sed

constitutive_law = NeoHookean

####################################################
## Define function to enforce essential boundary conditions
####################################################

abs_path= os.path.join(sys.path[0])

interior_points = device_put(jnp.load(f'{abs_path}/interior-points.npy')*1.)

boundary_adjust_fn = lambda U: U*interior_points

####################################################
## Define external forces (boundary and traction forces)
####################################################

final_pressure_loading = None
g                      = 980 # gravity force (data is in cm)
body_force             = jnp.array([0.]*(2) + [-g])

####################################################
## Define material parameter boundaries
####################################################

params_lb = jnp.array([500000.]*2)
params_ub = jnp.array([1000000.]*2)

####################################################
## Sample parameters on log scale or uniform scale 
####################################################

log_sampling = True

####################################################
## Set number of parameters to sample at each epoch 
####################################################

epoch_size = 200





data/OnceClampedBeam/geometryData/constitutive_law.py
dodaltuin/soft-tissue-pignn
"""
File: constitutive_law.py
Author: David Dalton
Description: Various setups before training, e.g.constitutive law
"""

####################################################
## Package imports
####################################################

import jax.numpy as jnp
from jax import device_put
import os
import sys

####################################################
## Transformation functions for stabilising strain 
## energy - see Section 2.4.3 of paper for details
####################################################

Jmin = 0.001  # minimum value of J which we allow
Jtrans = 0.05 # point at which transformation function kicks in

Jdiff = Jtrans - Jmin

beta_1 = 1. / Jdiff
beta_2 = jnp.log(Jdiff) - Jtrans/Jdiff
beta_3 = Jmin

exp_fn = lambda J: jnp.exp(beta_1*J + beta_2) + beta_3
J_transformation_fn = lambda J: jnp.where(J > Jtrans, J, exp_fn(J))

I1_trans = 10 # point at which transformation function kicks in

tanh_fn = lambda I1: jnp.tanh(I1 - I1_trans) + I1_trans
I1_trans_fn = lambda I1: jnp.where(I1 < I1_trans, I1, tanh_fn(I1))

####################################################
## Define constitutive law 
####################################################

def NeoHookean(params, F, J, fibres=None):
    """NeoHookean strain energy density function

    Inputs:
    -----------
    params: jnp.array
       Material parameter vector
    F     : jnp.array
       Deformation gradient
    J     : jnp.array
       Determinant of F
    fibres: None
       Not used here as we assume isotropic material

    Returns:
    ----------
    sed: float
       Strain energy density
    """

    # extract Lame material parameters
    mu, lambda_ = params

    # right Cauchy-Green tensor
    C = jnp.matmul(F.T, F)

    # first invariant
    Ic = jnp.trace(C)

    # transform to stop very large values
    Ic = I1_trans_fn(Ic)

    lnJ = jnp.log(J)

    sed = (mu/2.)*(Ic - 3.) - mu*lnJ + (lambda_/2.)*(lnJ**2)

    return sed

constitutive_law = NeoHookean

####################################################
## Define function to enforce essential boundary conditions
####################################################

abs_path= os.path.join(sys.path[0])

interior_points = device_put(jnp.load(f'{abs_path}/interior-points.npy')*1.)

boundary_adjust_fn = lambda U: U*interior_points

####################################################
## Define external forces (boundary and traction forces)
####################################################

final_pressure_loading = None
g                      = 980 # gravity force (data is in cm)
body_force             = jnp.array([0.]*(2) + [-g])

####################################################
## Define material parameter boundaries
####################################################

params_lb = jnp.array([500000.]*2)
params_ub = jnp.array([1000000.]*2)

####################################################
## Sample parameters on log scale or uniform scale 
####################################################

log_sampling = True

####################################################
## Set number of parameters to sample at each epoch 
####################################################

epoch_size = 200





data/TwiceClampedBeam/geometryData/constitutive_law.py
dodaltuin/soft-tissue-pignn
"""
File: constitutive_law.py
Author: David Dalton
Description: Various setups before training, e.g. constitutive law
"""

####################################################
## Package imports
####################################################

import jax.numpy as jnp
from jax import device_put
import os
import sys

####################################################
## Transformation functions for stabilising strain 
## energy - see Section 2.4.3 of paper for details
####################################################

Jmin = 0.001  # minimum value of J which we allow
Jtrans = 0.05 # point at which transformation function kicks in

Jdiff = Jtrans - Jmin

beta_1 = 1. / Jdiff
beta_2 = jnp.log(Jdiff) - Jtrans/Jdiff
beta_3 = Jmin

exp_fn = lambda J: jnp.exp(beta_1*J + beta_2) + beta_3
J_transformation_fn = lambda J: jnp.where(J > Jtrans, J, exp_fn(J))

I1_trans = 10 # point at which transformation function kicks in

tanh_fn = lambda I1: jnp.tanh(I1 - I1_trans) + I1_trans
I1_trans_fn = lambda I1: jnp.where(I1 < I1_trans, I1, tanh_fn(I1))

####################################################
## Define constitutive law 
####################################################

def NeoHookean(params, F, J, fibres=None):
    """NeoHookean strain energy density function

    Inputs:
    -----------
    params: jnp.array
       Material parameter vector
    F     : jnp.array
       Deformation gradient
    J     : jnp.array
       Determinant of F
    fibres: None
       Not used here as we assume isotropic material

    Returns:
    ----------
    sed: float
       Strain energy density
    """

    # extract material parameters
    E, nu = params

    # convert to Lame parameters
    mu, lambda_ = (E/(2*(1 + nu))), (E*nu/((1 + nu)*(1 - 2*nu)))

    # right Cauchy-Green tensor
    C = jnp.matmul(F.T, F)

    # first invariant
    Ic = jnp.trace(C)

    # transform to stop very large values
    Ic = I1_trans_fn(Ic)

    lnJ = jnp.log(J)

    sed = (mu/2.)*(Ic - 3.) - mu*lnJ + (lambda_/2.)*(lnJ**2)

    return sed

constitutive_law = NeoHookean

####################################################
## Define function to enforce essential boundary conditions
####################################################

abs_path= os.path.join(sys.path[0])

interior_points = device_put(jnp.load(f'{abs_path}/interior-points.npy')*1.)
disp_add = device_put(jnp.load(f'{abs_path}/disp-add.npy')*1.)

boundary_adjust_fn = lambda U: (U*interior_points) + disp_add

####################################################
## Define external forces (boundary and traction forces)
####################################################

final_pressure_loading = None
body_force             = None

####################################################
## Define material parameter boundaries
####################################################

params_lb = jnp.array([1.,  0.1])
params_ub = jnp.array([25., 0.4])

####################################################
## Sample parameters on log scale or uniform scale 
####################################################

log_sampling = True

####################################################
## Set number of parameters to sample at each epoch 
####################################################

epoch_size = 200

















data/TwistingBeam/geometryData/constitutive_law.py
dodaltuin/soft-tissue-pignn
"""
File: constitutive_law.py
Author: David Dalton
Description: Various setups before training, e.g.constitutive law
"""

####################################################
## Package imports
####################################################

import jax.numpy as jnp
from jax import device_put
import os
import sys

####################################################
## Transformation functions for stabilising strain 
## energy - see Section 2.4.3 of paper for details
####################################################

Jmin = 0.001  # minimum value of J which we allow
Jtrans = 0.05 # point at which transformation function kicks in

Jdiff = Jtrans - Jmin

beta_1 = 1. / Jdiff
beta_2 = jnp.log(Jdiff) - Jtrans/Jdiff
beta_3 = Jmin

exp_fn = lambda J: jnp.exp(beta_1*J + beta_2) + beta_3
J_transformation_fn = lambda J: jnp.where(J > Jtrans, J, exp_fn(J))

I1_trans = 10 # point at which transformation function kicks in

tanh_fn = lambda I1: jnp.tanh(I1 - I1_trans) + I1_trans
I1_trans_fn = lambda I1: jnp.where(I1 < I1_trans, I1, tanh_fn(I1))

####################################################
## Define constitutive law 
####################################################

def NeoHookean(params, F, J, fibres=None):
    """NeoHookean strain energy density function

    Inputs:
    -----------
    params: jnp.array
       Material parameter vector
    F     : jnp.array
       Deformation gradient
    J     : jnp.array
       Determinant of F
    fibres: None
       Not used here as we assume isotropic material

    Returns:
    ----------
    sed: float
       Strain energy density
    """

    # extract material parameters
    E, nu = params

    # convert to Lame parameters
    mu, lambda_ = (E/(2*(1 + nu))), (E*nu/((1 + nu)*(1 - 2*nu)))

    # right Cauchy-Green tensor
    C = jnp.matmul(F.T, F)

    # first invariant
    Ic = jnp.trace(C)

    # transform to stop very large values
    Ic = I1_trans_fn(Ic)

    lnJ = jnp.log(J)

    sed = (mu/2.)*(Ic - 3.) - mu*lnJ + (lambda_/2.)*(lnJ**2)

    return sed

constitutive_law = NeoHookean

####################################################
## Define function to enforce essential boundary conditions
####################################################

abs_path= os.path.join(sys.path[0])

interior_points = device_put(jnp.load(f'{abs_path}/interior-points.npy')*1.)
disp_add = device_put(jnp.load(f'{abs_path}/disp-add.npy')*1.)

boundary_adjust_fn = lambda U: (U*interior_points) + disp_add

####################################################
## Define external forces (boundary and traction forces)
####################################################

final_pressure_loading = None
body_force             = None

####################################################
## Define material parameter boundaries
####################################################

params_lb = jnp.array([1.,  0.1])
params_ub = jnp.array([25., 0.4])

####################################################
## Sample parameters on log scale or uniform scale 
####################################################

log_sampling = True

####################################################
## Set number of parameters to sample at each epoch 
####################################################

epoch_size = 200

















data/TwistingBeamUnprocessed/geometryData/constitutive_law.py
dodaltuin/soft-tissue-pignn
"""
File: data_main.py
Author: David Dalton
Description: Data processing script

Processes data from raw format into augmented graph format, with node and edge features assigned.

For details on the augmented graph format, see Dalton, CMAME 2022 and: https://github.com/dodaltuin/passive-lv-gnn-emul

Example usage: the following command can be used to process the Liver data:

python -m data_main --mode="run_all" --data_dir="Liver"

Once data processing is complete, emulator training and validation can be performed from the main directory
"""

import data_process_utils as utils

from typing import Sequence

from absl import app
from absl import flags
from absl import logging

flags.DEFINE_enum('mode', 'run_all', ['generate_topology', 'generate_nodes', 'generate_edges', 'run_all'], help = 'Select which data processing function to run')
flags.DEFINE_string('data_dir', 'TwistingBeamUnprocessed', help='Path to directory where raw data is stored')
flags.DEFINE_string('existing_topology_dir', "None", help='Path to directory where raw data is stored')
flags.DEFINE_string('existing_stats_dir', "None", help='(optional) Path to directory where normalisation summary stats are saved')
flags.DEFINE_integer('n_nearest_nbrs', 5, lower_bound=1, help='number of nearest neighbours to consider when defining root-root topology')
flags.DEFINE_integer('n_leaves', 4, lower_bound=2, help='(approximate) number of leaves each root node should have')
flags.DEFINE_integer('min_root_nodes', 2, lower_bound=1, help='minimum number of nodes to have final virtual layer')
FLAGS = flags.FLAGS

# (optional) hard code list of decreasing integers specifying the number of virtual nodes to create at each layer. If set to None, it will be calculated automatically using the shell input variables "n_leaves" and "min_root_nodes" defined above
n_nodes_per_layer: Sequence[int] = None

def main(_):

    if not utils.os.path.isdir(FLAGS.data_dir):
        raise NotADirectoryError(f'No directory at: {FLAGS.data_dir}')

    if FLAGS.mode in ['run_all', 'generate_topology']:
        utils.generate_augmented_topology(FLAGS.data_dir,
                                          FLAGS.existing_topology_dir,
                                          FLAGS.n_nearest_nbrs,
                                          n_nodes_per_layer,
                                          FLAGS.n_leaves,
                                          FLAGS.min_root_nodes)

    if FLAGS.mode in ['run_all', 'generate_nodes']:
        utils.generate_augmented_nodes(FLAGS.data_dir)

    if FLAGS.mode in ['run_all', 'generate_edges']:
        utils.generate_edge_features(FLAGS.data_dir)

if __name__ == "__main__":
    app.run(main)



data/data_main.py
dodaltuin/soft-tissue-pignn
"""
File: data_process_utils.py
Author: David Dalton
Description: Utility functions to process raw simulation data
"""

import os

import numpy as np
from numpy import newaxis

import os
import shutil

import jax
jax.config.update('jax_platform_name', 'cpu')
jax.config.update('jax_debug_nans', True)
from jax import vmap
import jax.numpy as jnp

from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors

from functools import partial
from typing import Sequence

from absl import logging

# some strings printed to console to seperate different results sections
SECTION_SEPERATOR = '##############################'
SECTION_SEPERATOR2 = SECTION_SEPERATOR * 2

# we normalise shape coefficients with respect to the variance on the below indexed column
REFERENCE_SHAPE_COEFF_COLUMN = 1

####################################
## Augmented Graph Topology Generation
####################################

def calc_n_nodes_per_layer(n_real_nodes: int, root_leaf_factor: int, min_root_nodes: int) -> Sequence[int]:
    """Calculates number of nodes in each layer of the augmented graph

    Parameters:
    -----------
    n_real_nodes: int
        The number of real nodes in the original graph representation
    root_leaf_factor: int
        The number of leaves to collapse into one root node (on average)
    min_root_nodes: int
        The minimum number of allowable root nodes in a single layer

    Returns:
    ----------
    n_nodes_per_layer: Sequence[int]
        A list of integers giving the number of nodes in each layer of the augmented graph.
        The 0'th index is the number of real nodes. The remainder gives the number in each
        of the virtual node layers, which is referred to as the vector $\mathbf{n}$ in
        Algorithm 1 of the manuscript. Each value is found by recursively integer dividing
        the previous value by "root_leaf_factor", starting with the value "n_real_nodes".
        This continues until the resulting value is less than the threshold specified by
        "min_root_nodes", when the routine stops
    """

    # initialise list to hold count of number of nodes in each node layer
    nodes_count_per_layer = []
    nodes_count_per_layer.append(n_real_nodes)

    i = 1
    while True:

        # number of nodes in layer i is the number of nodes from the previous layer
        # (integer) divided by "root_leaf_factor"
        nodes_count_layer_i = nodes_count_per_layer[i-1] // root_leaf_factor

        # exit if we have reached the specificed min number of virtual nodes allowed in a layer
        if nodes_count_layer_i < min_root_nodes:
            break

        nodes_count_per_layer.append(nodes_count_layer_i)

        i += 1

    return nodes_count_per_layer


def intra_layer_topology(positions: np.array, nn_count: int, shift_value: int) -> np.array:
    """Calculates connectivity between a single layer of virtual nodes

    Parameters:
    -----------
    positions: np.array
        An array giving the euclidean coordiantes of the virtual nodes to be connected
    nn_count: int
        The number of nearest neigbours to consider when forming node connections
    shift_value: int
        Amount by which (local) node indices need to be adjusted to align with
        their indices from the global graph structure

    Returns:
    ----------
    intra_layer_topology: np.array
        np.array representing the topology betweent the virtual node whose coordinates are
        given in positions. The array consists of two columns, the first being the sender
        node indices, and the second the receiver nodes. Note how these indices are shifted
        by "shift_value" so that they align with the global graph strucutre. The connections
        are found by connecting each node to its "nn_count" neighbours among all other nodes
    """

    # find "nn_count" nearest neighbours of each point in "positions"
    nodes_count = positions.shape[0]
    nbrs = NearestNeighbors(n_neighbors=(nn_count+1), algorithm='ball_tree').fit(positions)

    # find the indices of these nearest neighbours
    _, indices = nbrs.kneighbors(positions)

    # sender indices are simply 0, 1, ..., nodes_count
    senders = np.repeat(np.arange(nodes_count), nn_count).reshape(-1,1)

    # receiver indices come from the KNN results above
    receivers = indices[:,1:].flatten().reshape(-1,1)

    # indices of both senders and receivers are shifted to make sure they make sense
    # with respect to the entire augmented graph
    senders += shift_value
    receivers += shift_value

    return np.concatenate((senders, receivers), axis=1)


def intra_layer_topology_final(positions: np.array, shift_value: int) -> np.array:
    """Calculates complete connectivity between the last layer of virtual nodes

    Parameters:
    -----------
    positions: np.array
        An array giving the euclidean coordiantes of the virtual nodes to be connected
    shift_value: int
        Amount by which (local) node indices need to be adjusted to align with
        their indices from the global graph structure

    Returns:
    ----------
    intra_layer_topology: np.array
        np.array representing the topology between the final layer of virtual nodes,
        whose coordinates are given in positions, in the same format as outputted by
        the function "intra_layer_topology" above. In the last layer of nodes however,
        the nodes are fully connected
    """

    # create a directed edge between each pair of nodes
    n_nodes= positions.shape[0]
    topology = np.array([[i,j] for i in range(n_nodes) for j in range(n_nodes) if i != j])

    # the indices of topology are shifted so that they
    # make sense with respect to the entire augmented graph
    return topology + shift_value

def inter_layer_topology(leafnode_indices: Sequence[int], rootnode_indices: Sequence[int]) -> np.array:
    """Calculates connectivity from leaf to root nodes

    Parameters:
    -----------
    leafnode_indices: Sequence[int]
        Indices of the leaf nodes
    rootnode_indices: Sequence[int]
        Indices of the root nodes

    Returns:
    ----------
    intra_layer_topology: np.array
        np.array with two columns representing the connections from leaf to root nodes.
        The first column is simply the indices of the leaf nodes, "leafnode_indices", and
        the second column gives the indices of the root nodes, "rootnode_indices"
    """

    # the leaf nodes become senders
    senders = leafnode_indices

    # the root nodes the receivers
    receivers = rootnode_indices

    return np.concatenate((senders, receivers), axis=1)

def calculate_virtual_nodes_and_edges(real_nodes: np.array,
                                      real_topology: np.array,
                                      n_nearest_nbrs: int,
                                      n_nodes_per_layer: Sequence[int],
                                      n_leaf_nodes: int,
                                      min_root_nodes: int):
    """Generate Augmented Topology (Algorithm 1 from the manuscript)

    Parameters:
    -----------
    Are as explained in "generate_augmented_topology()" below

    Returns:
    ----------
    tuple
       The first value of which is the augmented graph topology used for GNN
       message passing for all graphs under considieration. Some other variables
       are returned to assist in generating the virtual nodes for different graphs
       of the class of problem under consideration
    """

    logging.info(SECTION_SEPERATOR)
    logging.info('Calling calculate_virtual_nodes_and_edges()')

    # initialise list to keep track of the root node indices for each leaf node
    # across all virtual node layers
    kmeans_labels_list = []

    # calculate number of nodes in each virtual node layer
    n_real_nodes = real_nodes.shape[0]
    if n_nodes_per_layer is None:
        n_nodes_per_layer = calc_n_nodes_per_layer(n_real_nodes, n_leaf_nodes, min_root_nodes)

    n_total_nodes = sum(n_nodes_per_layer)
    n_node_layers = len(n_nodes_per_layer)
    n_virtual_node_layers = n_node_layers - 1
    logging.info(f'n_virtual_node_layers: {n_virtual_node_layers}')
    logging.info(f'n_nodes_per_layer: {n_nodes_per_layer}')

    # boolean array which labels each real and virtual node with its respective node layer index, 0, ... n_node_layers-1
    node_layer_labels_nested = [[i]*count_i for i, count_i in enumerate(n_nodes_per_layer)]
    node_layer_labels = np.array([item for sublist in node_layer_labels_nested for item in sublist])

    # arbitrary (0-based) indexing of all nodes (real and virtual)
    node_indices = np.arange(n_total_nodes).reshape(-1,1)

    # dimensionality of the data (2d for beam data, 3d for LV data)
    n_dim = real_nodes.shape[1]

    # initialise all nodes (real and augmented, or $V \cup \tilde(V)$)
    all_nodes = np.zeros((n_total_nodes, n_dim), dtype=np.float32)
    all_nodes[node_layer_labels == 0] = real_nodes

    # initialise augmented topology ($\tilde(E)$) as np array
    augmented_topology = real_topology.copy()

    for layer_num in range(1, n_node_layers):

        logging.info(f'Processing virtual node layer {layer_num}')

        # the leaf nodes of virtual node layer "layer_num" are the nodes from the previous layer "layer_num-1"
        leaf_node_indices = node_indices[node_layer_labels == (layer_num - 1)]
        leaf_nodes = all_nodes[node_layer_labels == (layer_num - 1)]

        # cluster 'leaf_nodes' into 'nodes_count_per_layer[layer_num]' clusters
        kmeans = KMeans(n_clusters = n_nodes_per_layer[layer_num], random_state=0).fit(leaf_nodes)

        # the cluster centers obtained via k-means become the root nodes of this layer
        root_nodes = kmeans.cluster_centers_

        # write these coords to the array holding all node coords, at the correct rows
        all_nodes[node_layer_labels == layer_num] = root_nodes

        # save list which labels each leaf node the index of its root node found using kmeans above
        kmeans_labels_list.append(kmeans.labels_)

        # need to shift all node indices so that they make sense with respect to the entire augmented graph
        shift_value = np.sum(node_layer_labels < (layer_num))
        root_node_indices = (kmeans.labels_ + shift_value).reshape(-1,1)

        # calculate connectivty between the leaf and root nodes
        leaf_root_topology = inter_layer_topology(leaf_node_indices, root_node_indices)

        # calculate connectivity between the root nodes using "n_nearest_nbrs"
        if layer_num < n_virtual_node_layers:
            root_root_topology = intra_layer_topology(root_nodes, n_nearest_nbrs, shift_value)
        else:
            logging.info('Generating fully connected final virtal node layer\n')
            root_root_topology = intra_layer_topology_final(root_nodes, shift_value)

        # concatenate new node connections onto existing topology
        for new_topology in [leaf_root_topology, root_root_topology]:
            augmented_topology = np.concatenate((augmented_topology, new_topology), axis=0)

    # for each directed edge e_{ij}, we require that there is a twin edge e_{ji}, to ensure
    # conservation of momentum in the message passing stage of the GNN. The below makes
    # sure that this is the case
    reverse_topology = np.concatenate((augmented_topology[:,1:2], augmented_topology[:,0:1]), axis=1)
    augmented_topology = np.concatenate((augmented_topology, reverse_topology), axis=0)
    # remove duplicates
    augmented_topology = np.unique(augmented_topology, axis=0)

    return augmented_topology, kmeans_labels_list, node_layer_labels, all_nodes


def generate_augmented_topology(data_dir: str,
                                existing_topology_dir: str,
                                n_nearest_nbrs: int,
                                n_nodes_per_layer: Sequence[int],
                                n_leaves: int,
                                min_root_nodes: int):
    """Generate augmented topology on a representative geometry and process results

    Parameters:
    -----------
    Are as explained in "generate_augmented_topology()" below

    data_dir: str
       The directory where the raw simulation data is stored

    existing_topology_dir: str
       If using the topology data not stored in /data_dir, but instead existing
       augmented topology data already computed, set the path here

    n_nearest_nbrs: int
       Number of nearest neighbours to consider

    n_nodes_per_layer: Sequence[int] (OPTIONAL)
       As described in "calc_n_nodes_per_layer" above

    If n_nodes_per_layer is None, the below two inputs are used to calculate
    n_nodes_per_layer automatically:

    n_leaves: int
        The number of leaves to collapse into one root node (on average)

    min_root_nodes: int
        The minimum number of allowable root nodes in a single layer


    Returns:
    ----------
    Nothing is returned, instead the augmented topology for a fixed representative
    geometry is computed as in Algorithm 1 of the manuscript, with all results processed
    and saved for use for generating the augmented graph representation of all other
    simulation inputs.
    """


    logging.info(SECTION_SEPERATOR2)
    logging.info('Calling generate_augmented_graph()')

    # copy augmented topology data from "existing_topology_dir" if it exists
    if existing_topology_dir != "None":

        if not os.path.isdir(existing_topology_dir):
            raise NotADirectoryError(f'No directory at: {existing_topology_dir}')

        new_topology_dir = f'{data_dir}/topologyData'
        logging.info(f'Copying "{existing_topology_dir}" to  "{new_topology_dir}"')
        shutil.copytree(existing_topology_dir, new_topology_dir, dirs_exist_ok=False)

        return

    logging.info(f'n_leaves: {n_leaves}')
    logging.info(f'n_nearest_nbrs: {n_nearest_nbrs}')
    logging.info(SECTION_SEPERATOR2 + '\n')

    # load real node coordinates of representative graph
    ref_coords = np.load(f'{data_dir}/geometryData/reference-coords.npy')
    logging.info(f'ref_coords.shape: {ref_coords.shape}')

    # load the topology of these nodes
    real_topology = np.load(f'{data_dir}/geometryData/real-node-topology.npy').astype(np.int32)
    logging.info(f'real_topology.shape: {real_topology.shape}\n')

    # calculate the virtual nodes and edges of the augmented graph using Algorithm 2 of the manuscript
    results_tuple = calculate_virtual_nodes_and_edges(ref_coords, real_topology, n_nearest_nbrs, n_nodes_per_layer, n_leaves, min_root_nodes)

    # split out results tuple into its individual arrays
    augmented_topology, kmeans_labels_list, node_layer_labels, ref_augmented_coords = results_tuple

    # Only half of the directed edges in the augmented graph need to be assigned an edge feature vector,
    # because we enforce symmetry in the messages passed on pairs of directed edges. See Algorithm 3 of the manuscript
    # The edges we assign features vectors to are those in the "sparse_topology" defined below
    sparse_topology_indices = augmented_topology[:,0] > augmented_topology[:,1]
    sparse_topology = augmented_topology[sparse_topology_indices]

    logging.info(f'augmented_topology.shape    : {augmented_topology.shape}')
    logging.info(f'sparse_topology.shape       : {sparse_topology.shape}')
    logging.info(f'node_layer_labels.shape     : {node_layer_labels.shape}')
    logging.info(f'np.unique(node_layer_labels): {np.unique(node_layer_labels)}')
    logging.info(f'rep_augmented_coords.shape  : {ref_augmented_coords.shape}')
    logging.info(f'len(kmeans_labels_list)     : {len(kmeans_labels_list)}')

    # convert kmeans_labels_list to np array for saving
    if len(kmeans_labels_list) > 1:
        kmeans_labels_list = np.array(kmeans_labels_list, dtype=object)
    else:
        kmeans_labels_list = np.array(kmeans_labels_list, dtype=np.int32)

    logging.info(f'Saving generate_augmented_graph() results\n')
    np.save(f'{data_dir}/geometryData/augmented-topology.npy', augmented_topology)
    np.save(f'{data_dir}/geometryData/sparse-topology.npy', sparse_topology)
    np.save(f'{data_dir}/geometryData/reference-augmented-coords.npy', ref_augmented_coords)
    np.save(f'{data_dir}/geometryData/kmeans-labels-list.npy', kmeans_labels_list)
    np.save(f'{data_dir}/geometryData/node-layer-labels.npy', node_layer_labels)


####################################
## Augmented Node Generation
####################################

def aggregate_leaf_nodes(leaf_nodes: jnp.array, labels_list: Sequence[int], n_root_nodes: int) -> jnp.array:
    """Computes augmented nodes for all graphs in given datasets

    Parameters:
    -----------
    leaf_nodes: jnp.array
        Array giving the coords / features of the leaf node
    labels_list: Sequence[int]
        List which labels each leaf node in "leaf_nodes" to the corresponding root
        node to which it is grouped
    n_root_nodes: int
        Number of root nodes that "leaf_nodes" are being clustered to

    Returns:
    ----------
    root_nodes: jnp.array
       Returns the coords (features) of the each root node, found by
       averaging the coords (features) of their corresponding leaf nodes
    """

    # root node co-ordinates / features are the mean of all their respective leaf nodes
    sum_values = jax.ops.segment_sum(leaf_nodes, labels_list, n_root_nodes)

    # cardinalities of each root node cluster
    norm_values = jax.ops.segment_sum(np.ones((leaf_nodes.shape[0],1)), labels_list, n_root_nodes)

    # divide sum by cardinality to get average value
    root_nodes = sum_values / norm_values

    return root_nodes


def compute_augmented_nodes(real_nodes: jnp.array, real_node_features: jnp.array, node_layer_labels: Sequence[int], kmeans_labels_list: Sequence[Sequence[int]]):
    """Computes augmented node coords and features for given input graph

    Parameters:
    -----------
    real_nodes: jnp.array
        Array giving the coords of the real nodes in the graph
    real_node_features: jnp.array
        Array giving the feature vectors of the real nodes in the graph
    node_layer_labels: Sequence[int] >= 0
        Labels each node (real and virtual) with the node layer it is part of,
        starting with 0 for real nodes, 1 for first layer of virtual nodes, and
        so on
    kmeans_labels_list: Sequence[Sequence[int]]
        List of lists, where each element $l$ is a list giving the indices of the cluster
        centre to which the leaf nodes from step $l$ of Algorithm 1 of the manuscript were
        assigned

    Returns:
    ----------
    augmented_node_features: jnp.array
       Returns the augmented node features, found by recursively averaging the features
       from the previous layer, with respect to the common node topology as implied by "kmeans_labels"
    """

    n_dim1 = real_nodes.shape[-1]
    n_dim2 = real_node_features.shape[-1]

    # initialise arrays to hold augmented node coords and features
    augmented_nodes = jnp.zeros((node_layer_labels.shape[0], n_dim1), dtype=jnp.float32)
    augmented_node_features = jnp.zeros((node_layer_labels.shape[0], n_dim2), dtype=jnp.float32)

    # 0th layer of augmented nodes are simply the real nodes
    augmented_nodes = augmented_nodes.at[node_layer_labels==0].set(real_nodes)
    augmented_node_features = augmented_node_features.at[node_layer_labels==0].set(real_node_features)

    # loop over each layer of virtual nodes
    for i, labels_list_i in enumerate(kmeans_labels_list):

        leaf_nodes_i = augmented_nodes[node_layer_labels==i]
        leaf_node_features_i = augmented_node_features[node_layer_labels==i]

        # number of root nodes in this layer
        n_root_nodes = np.sum(node_layer_labels == (i+1))

        # aggregate leaf node coords to find root node coords
        root_nodes = aggregate_leaf_nodes(leaf_nodes_i, labels_list_i, n_root_nodes)
        augmented_nodes = augmented_nodes.at[node_layer_labels == (i+1)].set(root_nodes)

        # aggregate leaf node features to find root node features
        root_node_features = aggregate_leaf_nodes(leaf_node_features_i, labels_list_i, n_root_nodes)
        augmented_node_features = augmented_node_features.at[node_layer_labels == (i+1)].set(root_node_features)

    return augmented_node_features


def generate_augmented_nodes(data_dir: str):
    """Computes augmented nodes for all graphs in given datasets

    Parameters:
    -----------
    data_dir: str
        Path to directory where the data is stored

    Returns:
    ----------
    No values are returned, instead the function "compute_all_augmented_nodes" defined above is called
    """

    logging.info(SECTION_SEPERATOR2)
    logging.info('Calling generate_augmented_nodes()')
    logging.info(SECTION_SEPERATOR2 + '\n')

    # real node co-ordinates of the reference geometry
    ref_real_nodes = np.load(f'{data_dir}/geometryData/reference-coords.npy')

    # real node features of the reference geometry
    real_node_features = np.load(f'{data_dir}/geometryData/real-node-features.npy')

    # augmented topology ($E$ in manuscript) calculated on representative graph using the function "generate_augmented_graph()" defined above
    augmented_topology = np.load(f'{data_dir}/geometryData/augmented-topology.npy')

    # augmented node co-ordinates of the reference geometry
    ref_augmented_coords = np.load(f'{data_dir}/geometryData/reference-augmented-coords.npy')

    # gives the root node label in the immedietely higher layer of virtual nodes, for all nodes but laster layer of virtual nodes
    kmeans_labels_list = np.load(f'{data_dir}/geometryData/kmeans-labels-list.npy', allow_pickle=True)#.astype(np.int32)

    # variable indicating which layer each augmented node belongs to
    node_layer_labels = np.load(f'{data_dir}/geometryData/node-layer-labels.npy')

    # compute features of augmented nodes
    augmented_node_features = compute_augmented_nodes(ref_real_nodes, real_node_features, node_layer_labels, kmeans_labels_list)

    logging.info(f'augmented_node_features.shape: {augmented_node_features.shape}\n')


    node_features_mean = augmented_node_features.mean(0)
    node_features_std  = augmented_node_features.std(0)

    logging.info(f'node_features_mean: {node_features_mean}')
    logging.info(f'node_features_std : {node_features_std}')

    augmented_node_features = (augmented_node_features- node_features_mean) / node_features_std

    logging.info(f'saving node features')
    np.save(f'{data_dir}/geometryData/augmented-node-features.npy', augmented_node_features)

    stats_save_dir = f'{data_dir}/normalisationStatistics'
    if not os.path.isdir(stats_save_dir): os.mkdir(stats_save_dir)

    jnp.save(f'{stats_save_dir}/node-features-mean.npy', node_features_mean)
    jnp.save(f'{stats_save_dir}/node-features-std.npy' , node_features_std)


####################################
## Edge Feature Generation
####################################

def nodes_to_edge_features(nodes: jnp.array, senders: Sequence[int], receivers: Sequence[int]) -> jnp.array:
    """Computes edge features for all edges in graph simultaneously

    Parameters:
    -----------
    nodes: jnp.array
        Array of coords of all nodes (real and augmented) in graph
    senders: Sequence[int]
        Indices of sender nodes
    receivers: Sequence[int]
        Indices of receiver nodes

    Returns:
    ----------
    edge_features: jnp.array
        Array where each row gives the edge feature vector of the corresponding edge.
        The edge feature gives the relative position of the receiver node with respect
        to the sender node, and the distance between the two nodes

    """

    n_dim = nodes.shape[-1]

    # array of relative differences between senders and receiver nodes
    node_diff = nodes[receivers,:n_dim] - nodes[senders,:n_dim]

    # euclidean norm of each relative node differences (reshaped to 2d array)
    node_distance = jnp.sqrt((node_diff**2).sum(axis=1)).reshape(node_diff.shape[0],1)

    # edge features assigned to each edge is the concatenation of the above two arrays
    edge_features = jnp.concatenate((node_diff, node_distance), axis=1)

    return edge_features

def generate_edge_features(data_dir: str):
    """Computes edge features for reference geometry

    Parameters:
    -----------
    data_dir: str
        Path to directory where the data is stored

    Returns:
    ----------
    No values are returned, but the computed edge features are saved to data_dir/geometryData
    """

    logging.info(SECTION_SEPERATOR2)
    logging.info('Calling generate_edge_features()')
    logging.info(SECTION_SEPERATOR2 + '\n')

    # load "sparse" topology - one edge from each pair of "twin" directed edges in full augmented topology
    sparse_topology = jnp.load(f'{data_dir}/geometryData/sparse-topology.npy')
    logging.info(f'sparse_topology.shape: {sparse_topology.shape}\n')

    # load representative augmented coords
    ref_augmented_coords = np.load(f'{data_dir}/geometryData/reference-augmented-coords.npy')

    edge_features = nodes_to_edge_features(ref_augmented_coords, sparse_topology[:,0], sparse_topology[:,1])

    logging.info(f'edge_features.shape: {edge_features.shape}')

    edge_features_mean = edge_features.mean(0)
    edge_features_std  = edge_features.std(0)

    logging.info(f'edge_features_mean: {edge_features_mean}')
    logging.info(f'edge_features_std : {edge_features_std}')

    edge_features = (edge_features - edge_features_mean) / edge_features_std

    logging.info(f'saving edge features')
    jnp.save(f'{data_dir}/geometryData/edge-features.npy', edge_features)

    jnp.save(f'{data_dir}/normalisationStatistics/edge-features-mean.npy', edge_features_mean)
    jnp.save(f'{data_dir}/normalisationStatistics/edge-features-std.npy' , edge_features_std)


data/data_process_utils.py
dodaltuin/soft-tissue-pignn
"""
File: main.py
Author: David Dalton
Description: Wrapper script to call PI-GNN emulation
"""

#############################################################################
## Package imports
#############################################################################

from absl import app
from absl import flags
from utils_training import run_fns_dict

#############################################################################
## Set up shell input variables
#############################################################################

flags.DEFINE_enum(   'mode', 'train_and_evaluate', ['train', 'evaluate', 'train_and_evaluate'], help = 'Fit model to training data or evaluate on test data')
flags.DEFINE_integer('n_epochs', 1000, lower_bound = 1, help = 'Number of epochs to train the model for')
flags.DEFINE_string( 'data_path', 'LeftVentricle', help = 'Name of sub-directory in "/data" where model data is stored')
flags.DEFINE_float(  'lr', 1e-4, lower_bound=0, help='Learning rate for training the network')
flags.DEFINE_string( 'trained_params_dir', "None", help='Path to directory with pre-trained network parameters')
flags.DEFINE_integer('K', 5, lower_bound=1, help='Number of message passing steps to perform')
flags.DEFINE_string( 'dir_label', '', help='Optional label to append to end of results save directory')
FLAGS = flags.FLAGS

def main(_):

    for run_fn in run_fns_dict[FLAGS.mode]:

        run_fn(FLAGS.data_path,
               FLAGS.K,
               FLAGS.n_epochs,
               FLAGS.lr,
               FLAGS.trained_params_dir,
               FLAGS.dir_label
               )

if __name__ == "__main__":
    app.run(main)


main.py
dodaltuin/soft-tissue-pignn
"""
File: models.py
Author: David Dalton
Description: Implements PrimalGraphEmulator GNN Architecture
"""

import jax
import jax.numpy as jnp
import jax.tree_util as tree

from flax import linen as nn

from typing import Sequence, Callable

DTYPE = jnp.float32

class FlaxMLP(nn.Module):
  """Implements an MLP in Flax"""
  features: Sequence[int]
  layer_norm: bool

  @nn.compact
  def __call__(self, inputs):
    x = inputs
    for i, feat in enumerate(self.features):
      x = nn.Dense(feat, dtype=DTYPE)(x)
      if i != len(self.features) - 1:
        x = nn.celu(x)
    if self.layer_norm:
        x = nn.LayerNorm()(x)
    return x

def make_mlp(features: Sequence[int]):
    """Makes standard MLP

    With hidden layers defined by features
    """
    def update_fn(inputs):
        return FlaxMLP(features, False)(inputs)
    return update_fn

def make_layernorm_mlp(features: Sequence[int]):
    """Makes MLP followed by LayerNorm

    With hidden layers specified by features
    """
    def update_fn(inputs):
        return FlaxMLP(features, True)(inputs)
    return update_fn

def aggregate_incoming_messages(messages: jnp.ndarray, receivers: Sequence[int], n_nodes: int):
    """Sum aggregates incoming messages to each node

    Performs the sum over incoming messages $\sum_{j \in \mathcal{N}_i} m_{ij}^k$
    from the processor stage of Algorithm 2 of the manuscript, for all nodes $i$ similtaneously
    """
    return jax.ops.segment_sum(messages, receivers, n_nodes)

def MessagePassingStep(node_update_fn: FlaxMLP, edge_update_fn: FlaxMLP, senders: Sequence[int], receivers: Sequence[int], n_real_nodes: int):
    """Returns function to perform one message passing step

    Function _ApplyMessagePassingStep performs one step of message passing $k$ as
    in the for loop in Algorithm 2 of the manuscript.
    """

    def _ApplyMessagePassingStep(Vold: jnp.ndarray, Eold: jnp.ndarray):

        # calculate messages along each directed edge with an edge feature vector assigned
        messages = edge_update_fn(jnp.hstack((Eold, Vold[receivers], Vold[senders])))

        # aggregate incoming messages m_{ij} from nodes i to j where i > j
        received_messages_ij = aggregate_incoming_messages(messages, receivers, n_real_nodes)

        # aggregate incoming messages m_{ij} from nodes i to j where i < j
        # m_{ij} = -m_{ji} where i < j (momentum conservation property of the message passing)
        received_messages_ji = aggregate_incoming_messages(-messages, senders, n_real_nodes)

        # concatenate node representation with incoming messages and then update node representation
        V = node_update_fn(jnp.hstack((Vold, received_messages_ij + received_messages_ji)))

        # return updated node and edge representations with residual connection
        return Vold + V, Eold + messages

    return _ApplyMessagePassingStep


class PrimalGraphEmulator(nn.Module):
    """PrimalGraphEmulator (varying geometry data)"""
    mlp_features: Sequence[int]
    latent_size: Sequence[int]
    K: int
    receivers: Sequence[int]
    senders: Sequence[int]
    n_total_nodes: int
    output_dim: int
    real_node_indices: Sequence[bool]
    boundary_adjust_fn: Sequence[bool] = None

    @nn.compact
    def __call__(self, V: jnp.ndarray, E: jnp.ndarray, theta: jnp.ndarray, sow_latents=False):
        """Implements Algorithm 2 of the manuscript: forward pass of PrimalGraphEmulator

        Inputs:
        ---------
        V: jnp.ndarray
            Array giving feature vectors of each node (real and virtual)
        E: jnp.ndarray
            Array giving feature vectors of each edge
        theta: jnp.ndarray
            Vector of global graph parameters

        Outputs:
        ---------
        U: jnp.ndarray
            Array of displacement predictions for each real node in V
        """

        ## Initialise internal MLPs:
        # 3 encoder MLPs
        node_encode_mlp=make_layernorm_mlp(self.mlp_features + self.latent_size)
        edge_encode_mlp=make_layernorm_mlp(self.mlp_features + self.latent_size)
        theta_encode_mlp=make_layernorm_mlp(self.mlp_features + self.latent_size)

        # 2K processor MLPs
        message_passing_blocks = [MessagePassingStep(make_layernorm_mlp(self.mlp_features + self.latent_size),
                                                     make_layernorm_mlp(self.mlp_features + self.latent_size),
                                                     self.senders, self.receivers, self.n_total_nodes) for i in range(self.K)]

        # D decoder MLPs
        node_decode_mlps = [make_mlp(self.mlp_features + (1,)) for i in range(self.output_dim[0])]

        ## Encoder:
        V = node_encode_mlp(V)
        E = edge_encode_mlp(E)

        ## Processor:
        # perform K rounds of message passing
        for message_pass_block_i in message_passing_blocks:
            V, E = message_pass_block_i(V, E)

        # aggregate incoming messages to each node
        incoming_messages = aggregate_incoming_messages(E, self.receivers, self.n_total_nodes)

        # final local learned representation is a concatenation of vector embedding and incoming messages
        z_local = nn.LayerNorm()(jnp.hstack((V, incoming_messages))[self.real_node_indices])

        # used for rapid evaluations of decoder for fixed geometry
        if sow_latents: return z_local

        ## Decoder:
        # encode global parameters theta
        z_theta = theta_encode_mlp(theta)

        # tile global parameter embeddings (z_theta) to each individual real node
        z_theta_array = jnp.tile(z_theta, (z_local.shape[0], 1))

        # final learned representations are comprised of (z_theta, z_local)
        final_learned_representations = jnp.hstack((z_theta_array, z_local))

        # make prediction for forward displacement using different decoder mlp for each dimension
        individual_mlp_predictions = [decode_mlp(final_learned_representations) for decode_mlp in node_decode_mlps]

        # concatenate the predictions of each individual decoder mlp
        Upred = jnp.hstack(individual_mlp_predictions)

        # adjust predictions to account for displacement boundary conditions
        if self.boundary_adjust_fn is not None:
            Upred = self.boundary_adjust_fn(Upred)

        # return displacment prediction array
        return Upred


class PrimalGraphEmulatorDecoder(nn.Module):
    """PrimalGraphEmulator (just decoder state)"""
    mlp_features: Sequence[int]
    output_dim: Sequence[int]
    n_real_nodes: int
    latent_nodal_values: jnp.ndarray
    theta_encode_mlp_fn: Callable
    boundary_adjust_fn: Sequence[float] = None

    @nn.compact
    def __call__(self, theta: jnp.ndarray):
        """Implements Decoder stage PrimalGraphEmulator

        Input:
        ---------
        theta: jnp.array
            Vector giving the global parameters $\theta$ for the fixed geom
            being modelled

        Output:
        ---------
        U: jnp.ndarray
            Array of displacement predictiosn for each real node
        """


        # initialise node-decode MLP
        node_decode_mlps = [make_mlp(self.mlp_features + (1,)) for i in range(self.output_dim[0])]

        # embed theta to higher dim space using pre-trained theta_encode_mlp
        z_theta = self.theta_encode_mlp_fn(theta)

        # tile global parameter embeddings (z_theta) to each individual real finite-element node
        z_theta_array = jnp.tile(z_theta, (self.n_real_nodes,1))

        # final learned representations are comprised of (z_theta, z_local)
        final_learned_representations = jnp.hstack((z_theta_array, self.latent_nodal_values))

        # make prediction for forward displacement using different decoder mlp for each dimension
        individual_mlp_predictions = [decode_mlp(final_learned_representations) for decode_mlp in node_decode_mlps]

        # concatenate the predictions of each individual decoder mlp
        Upred = jnp.hstack(individual_mlp_predictions)

        # adjust predictions to account for displacement boundary conditions
        if self.boundary_adjust_fn is not None:
            Upred = self.boundary_adjust_fn(Upred)

        # return displacment prediction array
        return Upred

models.py
dodaltuin/soft-tissue-pignn
"""
File: utils.py
Author: David Dalton
Description: Utility functions for initialising PI-GNN emulators and saving results
"""

import os

import pathlib
import pickle

from jax import random
import jax.numpy as jnp
import flax
from flax.core import freeze, unfreeze

from typing import Sequence

import models

def create_config_dict(K: int, n_epochs: int, lr: float, output_dim: int, local_embed_dim: int, mlp_features: Sequence[int], rng_seed: int):
    """Creates dictionary of configuration details for the GNN emulator"""

    return {'K': K,
            'n_train_epochs': n_epochs,
            'learning_rate': lr,
            'output_dim': output_dim,
            'local_embedding_dim': local_embed_dim,
            'mlp_features': mlp_features,
            'rng_seed': rng_seed
            }

def create_savedir(data_path: str, K: int, n_epochs: int, lr: float, dir_label: str, local_embedding_dim: int, mlp_width: int, mlp_depth: int, rng_seed: int):
    """Create directory where emulation results are saved

    The emulator's configuration details are written to the directory name for ease of reference
    """

    save_dir = f'emulationResults/{data_path}/pge_K{K}_Ntrain{n_epochs}_lr{lr:.1e}_M{local_embedding_dim}_FCNN_W{mlp_width}_FCNN_D{mlp_depth}_RNG{rng_seed}{dir_label}/'
    if not os.path.isdir(save_dir): os.makedirs(save_dir)

    # also create a directory to save log files
    log_dir = f'{save_dir}/logFiles'
    if not os.path.isdir(log_dir): os.makedirs(log_dir)

    return save_dir


def load_trained_params(params_save_dir, params_filename="trainedNetworkParams.pkl"):
    """Load trained parameters of GNN emulator from params_save_dir"""

    params_filename_full = params_save_dir + params_filename
    if not os.path.isfile(params_filename_full):
        raise FileNotFoundError(f'No file at: {params_filename_full}')

    with pathlib.Path(params_filename_full).open('rb') as fp:
        params_load = pickle.load(fp)

    return params_load


def gen_zero_params_gnn(model, params_randL):
    node_decode_first_mlp_index = model.K*2 + 3
    node_decode_last_mlp_index = node_decode_first_mlp_index + model.output_dim[0]

    mlp_depth = len(model.mlp_features)
    index = node_decode_first_mlp_index

    final_weights_layer_decoder_rand = params_randL['params'][f'FlaxMLP_{index}'][f'Dense_{mlp_depth}']['kernel']
    final_weights_layer_decoder_zero = jnp.zeros_like(final_weights_layer_decoder_rand)

    params_zero = unfreeze(params_randL)

    for index in range(node_decode_first_mlp_index, node_decode_last_mlp_index):
        params_zero['params'][f'FlaxMLP_{index}'][f'Dense_{mlp_depth}']['kernel'] = final_weights_layer_decoder_zero

    return freeze(params_zero)


def initialise_network_params(data_generator, ref_geom, model, trained_params_dir: str, rng_seed: int):
    """Initialise the parameters of the GNN emulator

    If initialising from scratch, use the ".init" method from Flax

    If initialising from earlier training results, simply read these parameters
    from trained_params_dir
    """

    if trained_params_dir == "None":
        key = random.PRNGKey(rng_seed)
        theta_init, _ = data_generator.get_data(0)
        V_init = ref_geom._node_features
        E_init = ref_geom._edge_features

        params = model.init(key, V_init, E_init, theta_init)
        return params
    else:
        trained_params = load_trained_params(trained_params_dir)
        return trained_params


def init_emulator_full(config_dict: dict, data_generator, trained_params_dir: str, ref_geom):
    """Initialise GNN emulator (varying geometry data)

    Initialises GNN architecture and trainable paramters for prediction of varying LV geom data

    If trained_params_dir is "None", the parameters are initialised randomly
    If trained_params_dir is a directory path, pre-trained parameters are read from there
    """

    # initialise GNN architecture based on configuration details
    emulator =  models.PrimalGraphEmulator(mlp_features=config_dict['mlp_features'],
                                           latent_size=[config_dict['local_embedding_dim']],
                                           K = config_dict['K'],
                                           receivers = ref_geom._receivers,
                                           senders = ref_geom._senders,
                                           n_total_nodes= ref_geom._n_total_nodes,
                                           output_dim= [config_dict['output_dim']],
                                           real_node_indices = ref_geom._real_node_indices,
                                           boundary_adjust_fn = ref_geom.boundary_adjust_fn)

    # initialise trainable emulator parameters (either randomly or read from trained_params_dir)
    params = initialise_network_params(data_generator, ref_geom, emulator, trained_params_dir, config_dict['rng_seed'])

    return emulator, params


def init_emulator_decoder(model, params: dict, data_generator, ref_geom, results_save_dir: str, trained_params_dir: str):
    """Initialises fixed-geom emulator network parameters

    Initialises the parameters of a GNN emulator for fixed LV geometry data,
    where the pre-trained parameters come from a varying geometry emulator

    In this case, the message passing stage for the fixed geometry must
    be pre-compututed before the emulator (DeepGraphEmulatorFixedGeom) is
    initialised
    """

    # extract the node features (V), edge features (E), global parameters (theta)
    # and global shape embedding coefficients (z_global)
    #V, E, theta, z_global, _, _ = data_loader.return_index_0()
    theta_init = None
    V_init = ref_geom._node_features
    E_init = ref_geom._edge_features
    z_global_init = None

    # sow_latents=True returns the local learned reprentation from the message
    # passing stage or each node in the geometry
    latent_nodal_values = model.apply(params, V_init, E_init, theta_init, sow_latents=True)

    # Flax names the internal MLPs of the GNN 'FlaxMLP_i', where i ranges over the
    # number of MLPs in the network, and is numbered in order of initialisation
    # There are 2 (the two encode MLPs) + K*2 (two MLPs for each message passing step)
    # + 1 (the theta encode MLP) + D decoder MLPs = 3 + K*2 + D MLPs in the DeepGraphEmulator
    # GNN architecture. We want the index of the theta encoder MLP for use in the fixed
    # geometry emulator, which was the (2*K + 3)th MLP to be initialised. Therefore its
    # index is given as follows (Python uses 0-based indexing)
    theta_encode_mlp_index = model.K*2 + 2

    # extract theta encoder MLP params from params dictionary
    theta_encode_mlp_params = params['params'][f'FlaxMLP_{theta_encode_mlp_index}']
    theta_encode_mlp_params_dict = {'params': theta_encode_mlp_params.unfreeze()}
    theta_encode_mlp_params_dict = flax.core.frozen_dict.freeze(theta_encode_mlp_params_dict)

    # define a function to map theta to z_theta given value for theta encode MLP parameters
    theta_encode_mlp = models.FlaxMLP(model.mlp_features + model.latent_size, True)
    theta_encode_mlp_fn = lambda x: theta_encode_mlp.apply(theta_encode_mlp_params_dict, x)

    # initialise the fixed geometry (fg) emulator
    model_fg = models.PrimalGraphEmulatorDecoder(mlp_features = model.mlp_features,
                                                 output_dim= [ref_geom._output_dim],
                                                 n_real_nodes= ref_geom._n_real_nodes,
                                                 latent_nodal_values= latent_nodal_values,
                                                 theta_encode_mlp_fn=theta_encode_mlp_fn,
                                                 boundary_adjust_fn = ref_geom.boundary_adjust_fn)

    # index for first decoder MLP is one after the theta_encode_mlp
    node_decode_first_mlp_index = theta_encode_mlp_index + 1
    node_decode_last_mlp_index = node_decode_first_mlp_index + model.output_dim[0]

    # extract parameters for "D" node decode MLPs to a list
    decoder_mlps_params_list = [params['params'][f'FlaxMLP_{index}'] for index in range(node_decode_first_mlp_index, node_decode_last_mlp_index)]

    # convert parameters list to a frozen_dict suitable for use with Flax
    decoder_mlps_params_dict = {'params': {f'FlaxMLP_{i}':params_i for i, params_i in enumerate(decoder_mlps_params_list)}}
    params_fg = flax.core.frozen_dict.freeze(decoder_mlps_params_dict)

    return model_fg, params_fg


def initialise_emulator(emulator_config_dict, data_generator, results_save_dir, ref_geom_data, fix_message_layers=False, trained_params_dir="None"):
    """Initialises GNN emulator

    Initialises GNN emulator and parameters, given specifications given in emulator_config_dict

    If fixed_message_layers=False, returns instance of models.PrimalGraphEmulator
    If fixed_message_layers=True, returns instance of models.PrimalGraphEmulatorDecoder
    """

    # check input conflict
    if fix_message_layers:
        assert trained_params_dir != "None", \
               'If fixing message passing layers, must initialise from pre-trained network parameters'

    # initialise varying geometry emulator (models.PrimalraphEmulator) and parameters
    emulator, params = init_emulator_full(emulator_config_dict, data_generator, trained_params_dir, ref_geom_data)
    emulator_pred_fn = lambda p, theta_norm: emulator.apply(p, ref_geom_data._node_features, ref_geom_data._edge_features, theta_norm)

    # if fixing message passing layers, use the above model and parameters to initialise the decoder stage of the emulator by pre-computing the message passing stage of the emulation architecture
    if fix_message_layers:
        emulator, params = init_emulator_decoder(emulator, params, data_generator, ref_geom_data, results_save_dir, trained_params_dir)
        emulator_pred_fn = lambda p, theta_norm: emulator.apply(p, theta_norm)

    return emulator_pred_fn, params, emulator


utils.py
dodaltuin/soft-tissue-pignn
"""
File: utils_data.py
Author: David Dalton
Description: data loading utility classes for model and simulation data
"""

import numpy as np
import jax.numpy as jnp
from jax import device_put, random

from scipy.stats.qmc import LatinHypercube, scale as qmc_scale

import os
import sys

from typing import Sequence

import utils_potential_energy as utils_pe

class ExternalForces:
    """Class for storing any external forces that act on the geometry"""

    def __init__(self, data_path: str):

        geometry_data_dir = f'data/{data_path}/topologyData'

        sys.path.insert(0, geometry_data_dir)
        from constitutive_law import final_pressure_loading, body_force

        self.body_force    = body_force             # jnp array (vector)
        self.surface_force = final_pressure_loading # float


class DataGenerator:
    """Class for generating input theta data points at which the PI-GNN emulator is trained"""

    def __init__(self, data_path: str, lhs_seed: int = 101132, sampler_seed: int = 42, shuffle_seed: int = 420):
        """
        Parameters
         ----------
        data_path: str
               Name of the subdirectory within "/data" where the data is stored
        *_seed: int
               Random seeds for data sampling/shuffling
        """

        geometry_data_dir = f'data/{data_path}/geometryData'
        stats_dir         = f'data/{data_path}/normalisationStatistics'

        sys.path.insert(0, geometry_data_dir)
        from constitutive_law import params_lb, params_ub, epoch_size, log_sampling

        # can optionally specify to sample theta inputs on log scale between upper and lower bounds
        if log_sampling:
            self.transform_fn = jnp.log
            self.transform_inv = jnp.exp
        else:
            identity_fn = lambda x: x
            self.transform_fn  = identity_fn
            self.transform_inv = identity_fn

        self.params_lb = self.transform_fn(params_lb)
        self.params_ub = self.transform_fn(params_ub)

        # initialse sampler and shuffler random seeds
        self.lhs_seed = lhs_seed
        self.sampler_key = random.PRNGKey(sampler_seed)
        self.shuffle_key = random.PRNGKey(shuffle_seed)

        self.n_params = len(self.params_lb)

        # array of data point indices that can be iterated over during each epoch
        self.epoch_size = epoch_size
        self.epoch_indices = jnp.arange(self.epoch_size)

        # generate input data points using LHS sample
        self.generate_lhs_points(stats_dir)

    def generate_lhs_points(self, stats_dir):
        """Generate input theta points using Latin HyperCube sampling"""

        sampler = LatinHypercube(d=len(self.params_lb), seed=self.lhs_seed)

        hypercube_samples = sampler.random(n=self.epoch_size)

        theta = self.transform_inv(qmc_scale(hypercube_samples, self.params_lb, self.params_ub))

        self.theta = device_put(theta)

        self.theta_mean = self.theta.mean(0)
        self.theta_std = self.theta.std(0)

        np.save(f'{stats_dir}/theta-mean.npy', self.theta_mean)
        np.save(f'{stats_dir}/theta-std.npy', self.theta_std)

        self.theta_norm = (self.theta - self.theta_mean) / self.theta_std

    def shuffle_epoch_indices(self):
        """Shuffles the order in which the dataset is cycled through

        This is called at the start of each training epoch to randomise the order in which the input data points are seen
        """

        self.shuffle_key, key = random.split(self.shuffle_key)
        self.epoch_indices = random.choice(key, self.epoch_indices, shape=(self.epoch_size,), replace=False)

    def resample_input_points(self):
        """Resample the input points over which the emulator is trained"""

        # reset sampler key
        self.sampler_key, key = random.split(self.sampler_key)

        # uniform random sampling between the lower and upper bounds
        samples = random.uniform(key, shape=(self.epoch_size, self.n_params), minval=self.params_lb, maxval=self.params_ub)

        # save results in original space and normalised results
        self.theta = device_put(self.transform_inv(samples))
        self.theta_norm = device_put((self.theta - self.theta_mean) / self.theta_std)

    def get_data(self, data_idx):
        """Returns input global graph values for specified data point (normalised and unnormalised)"""

        return self.theta_norm[data_idx], self.theta[data_idx]


class ReferenceGeometry:
    """Class for storing geometric/topological data for reference body"""

    def __init__(self, data_path: str, logging=None):
        """
        Parameters
        ----------
        data_path: str
            Name of the subdirectory within "/data" where the data is stored
        logging
            logging module to write information to
        """

        full_data_path = f'data/{data_path}'
        if not os.path.isdir(full_data_path):
            raise NotADirectoryError(f'No directory at: {full_data_path}')

        geometry_data_dir = f'{full_data_path}/geometryData'
        stats_dir         = f'{full_data_path}/normalisationStatistics'

        # load geometry coords in reference configuration
        self.ref_coords = device_put(jnp.load(f'{geometry_data_dir}/reference-coords.npy'))
        self._n_real_nodes, self._output_dim = self.ref_coords.shape

        sys.path.insert(0, geometry_data_dir)
        from constitutive_law import constitutive_law, J_transformation_fn, boundary_adjust_fn

        # constitutive law for the material
        self.constitutive_law = constitutive_law

        # function to transform J to prevent negative values
        self.Jtransform = J_transformation_fn

        self.boundary_adjust_fn = boundary_adjust_fn

        # load mesh finite elements
        self.elements = device_put(jnp.load(f'{geometry_data_dir}/elements.npy'))

        # compute volumes of the finite elements
        self.element_vols = utils_pe.compute_vol(self.elements, self.ref_coords).reshape(-1,1)

        # load mesh topology
        sparse_topology =  jnp.load(f'{geometry_data_dir}/sparse-topology.npy').astype(jnp.int32)
        self._senders = sparse_topology[:,0]
        self._receivers = sparse_topology[:,-1]

        # node features
        self._node_features = device_put(jnp.load(f'{geometry_data_dir}/augmented-node-features.npy').astype('float32'))

        # edge features
        self._edge_features = device_put(jnp.load(f'{geometry_data_dir}/edge-features.npy').astype(np.float32))

        # fibre field (for non-isotropic materials), defined element-wise
        fibre_field_file =   f'{geometry_data_dir}/fibre-field-elementwise.npy'
        if os.path.exists(fibre_field_file):
            self._fibre_field = device_put(np.load(fibre_field_file))
        else:
            # for isotropic materials the fibre field is set to None
            self._fibre_field = [None]*self.elements.shape[0]

        # for models involving a surface force, the below file indicates the nodes and surface facets which lie on the surface where the force is applied
        pressure_surface_indicator_file = f'{geometry_data_dir}/pressure-surface-facets.npy'
        if os.path.exists(pressure_surface_indicator_file):

            pressure_surface_facets = jnp.array((jnp.load(pressure_surface_indicator_file)))

            # compute areas and normal vectors of each triangular facet on the pressure surface ($\partial \Omega^\sigma$ in Eq. 1 of the manuscript)
            area_facets, N_facets = utils_pe.compute_area_N(pressure_surface_facets, self.ref_coords, self.elements)

            # save results
            self.tri_surface_area_normals    = device_put(area_facets.reshape(-1,1) * N_facets)
            self.tri_surface_element_indices = device_put(pressure_surface_facets[:,0])
            self.tri_surface_node_indices    = device_put(pressure_surface_facets[:,1:])

        # label each node as layer 0 (real nodes), layer 1 (first layer of virtual nodes), etc
        node_layer_labels = np.load(f'{geometry_data_dir}/node-layer-labels.npy')
        self._real_node_indices = (node_layer_labels == 0)

        self._n_total_nodes = self._node_features.shape[0]
        self._n_edges= self._edge_features.shape[0]
        self._output_dim = self.ref_coords.shape[-1]

        if logging is not None:
            logging.info(f'{data_path} geometry summary:')
            logging.info(f'n_real_nodes  : {self._n_real_nodes}')
            logging.info(f'n_total_nodes : {self._n_total_nodes}')
            logging.info(f'n_edges       : {self._n_edges}')
            logging.info(f'elements.shape: {self.elements.shape}')
            logging.info(f'edges.shape   : {self._edge_features.shape}')
            logging.info(f'nodes.shape   : {self._node_features.shape}')
            if os.path.exists(pressure_surface_indicator_file):
                logging.info(f'tri_surface_area_normals.shape   : {self.tri_surface_area_normals.shape}\n')


class DataLoader:
    """Data loader for graph-formatted input-output simulation data with common, fixed topology"""

    def __init__(self, data_path: str, data_type: str):
        """
        Parameters
        ----------
        data_path: str
            Name of the subdirectory within "/data" where the data is stored
        data_type: str
            The type of data to be loaded, one of "train", "validation", "test"
        """

        full_data_path = f'data/{data_path}'
        if not os.path.isdir(full_data_path):
            raise NotADirectoryError(f'No directory at: {full_data_path}')

        sim_data_dir      = f'{full_data_path}/simulationData/{data_type}'
        stats_dir         = f'{full_data_path}/normalisationStatistics'
        geometry_data_dir = f'{full_data_path}/geometryData'

        sys.path.insert(0, geometry_data_dir)
        from constitutive_law import log_sampling

        # simulated displacements
        self._displacement = jnp.load(f'{sim_data_dir}/displacement.npy')

        # total potential energy values returned by the simulator
        self._pe_values = np.load(f'{sim_data_dir}/pe-values.npy')

        # global graph parameters for each simulation (generally material parameters)
        self._theta = device_put(jnp.load(f'{sim_data_dir}/theta.npy'))

        # normalisation statistics for theta
        self.theta_mean = jnp.load(f'{stats_dir}/theta-mean.npy')
        self.theta_std = jnp.load(f'{stats_dir}/theta-std.npy')

        # normalised theta values for input to GNN
        theta_norm = (((self._theta)) - self.theta_mean ) / self.theta_std
        self._theta_norm = device_put(theta_norm)

        # array of data point indices that can be iterated over during each epoch
        self._data_size = self._theta.shape[0]
        self._epoch_indices = np.arange(self._data_size)

    def get_graph(self, data_idx):
        """Returns input/output values for specified data point and places on GPU"""

        return device_put(self._theta_norm[data_idx:(data_idx+1)]), device_put(self._displacement[data_idx]), device_put(self._theta[data_idx]), self._pe_values[data_idx]

    def get_data(self, data_idx):
        """Returns input global graph values for specified data point (normalised and unnormalised)"""

        return self._theta_norm[data_idx], self._theta[data_idx]

    def return_index_0(self):
        """Returns input/output data ([V, E, theta]/ U) for first data point and places on GPU

        This method is used when initialising the parameters of the varying geometry GNN emulator
        """

        return device_put(self._nodes[0]), device_put(self._edges[0]), device_put(self._theta_vals[0:1]), device_put(self._shape_coeffs[0]), device_put(self._displacement[0]), device_put(self._theta[0])

    def shuffle_epoch_indices(self, seed_idx):
        """Shuffles the order in which the dataset is cycled through"""

        np.random.seed(seed_idx)
        self._epoch_indices = np.random.choice(self._data_size, self._data_size, False)


utils_data.py
dodaltuin/soft-tissue-pignn
"""
File: utils_evaluation.py
Author: David Dalton
Description: Utility functions for evaluating emulator prediction accuracy
"""

import jax
import jax.numpy as jnp
import numpy as np
import shutil
import os

from utils_potential_energy import compute_def_gradient

def compute_I1(F):
    """Compute first invariant of the deformation gradient F"""
    return jnp.trace(jnp.matmul(F.T, F))

compute_I1_vmap = jax.vmap(compute_I1)

def compute_F_J_I1(displacement, ref_geom_data):
    """Compute deformation gradient F, J=det(F) and first invariant I1"""

    # coords in current configuration after displacement
    cur_coords = ref_geom_data.ref_coords + displacement

    F, J = compute_def_gradient(ref_geom_data.elements, ref_geom_data.ref_coords, cur_coords, ref_geom_data.Jtransform)

    I1 = compute_I1_vmap(F)

    return F, J, I1

compute_F_J_I1_vmap = jax.vmap(compute_F_J_I1, in_axes = [0] + [None])

def err_u(u, uhat=0.):
    """RMSE of nodal displacement predictions"""
    return (((u - uhat)**2).sum(-1))**.5

def rel_err_u(u, uhat):
    """Relative error of displacement predictions"""
    return err_u(u, uhat)/err_u(u)

# matrix norm
Frobenius_norm = lambda X: (X**2).sum([-1, -2])**.5

def err_F(F, Fhat):
    """Error in prediction of deformation gradient"""
    return 100.*Frobenius_norm(F-Fhat)/Frobenius_norm(F)

def err_norm(true, pred):
    return jnp.abs((true-pred)/true)*100.

def find_quantiles(arr, label, logging):
    quantiles = jnp.percentile(arr, jnp.array([0., 2.5, 25., 50., 75., 97.5, 100.]))
    logging.info(f'{label} quantiles: 0%:{quantiles[0]:.2e}, 2.5%:{quantiles[1]:.2e}, 25%:{quantiles[2]:.2e}, 50%:{quantiles[3]:.2e}, 75%:{quantiles[4]:.2e}, 97.5%:{quantiles[5]:.2e}, 100%:{quantiles[-1]:.2e}')
    return quantiles

def print_error_statistics(true_arrs, pred_arrs, save_dir, logging):
    """Prints prediction error statistics to console"""

    Utrue, PEtrue, Ftrue, Jtrue, I1true = true_arrs
    Upred, PEpred, Fpred, Jpred, I1pred = pred_arrs

    u_norms  = err_u(Utrue)
    u_errors = err_u(Utrue, Upred)
    PE_errors = err_norm(PEtrue, PEpred)
    F_errors = err_F(Ftrue, Fpred)
    J_errors = err_norm(Jtrue, Jpred)
    I1_errors = err_norm(I1true, I1pred)

    logging.info('Prediction error quantiles:')
    find_quantiles(u_norms,   'u norm', logging)
    find_quantiles(u_errors,  'err_u ', logging)
    find_quantiles(PE_errors, 'err_PE', logging)
    find_quantiles(F_errors,  'err_F ', logging)
    find_quantiles(J_errors,  'err_J ', logging)
    find_quantiles(I1_errors, 'err_I1', logging)

def find_quantile_indices(loss_array):
    """find the simulations from the test set on which the emulator achieved the median and worst prediction error respectively"""

    # median and worst loss values
    loss_quantiles = [50, 100]
    loss_quantile_values = np.percentile(loss_array, loss_quantiles)

    index_values = []

    # find the corresponding indices in the data set
    for loss_quantile_i in loss_quantile_values:
        index = np.argmin(np.abs(loss_array - loss_quantile_i))
        index_values.append(index)

    return index_values

utils_evaluation.py
dodaltuin/soft-tissue-pignn
"""
File: utils_potential_energy.py
Author: David Dalton - adjusted from Matlab code written by Hao Gao (https://github.com/HaoGao)
Description: Utility functions for calculating total potential energy
"""

import jax
import jax.numpy as jnp

from typing import Callable

def compute_area_N_facet(facet: jnp.ndarray, coords: jnp.ndarray, elements: jnp.ndarray):
    """Compute area and normal vector N for a triangular surface facet """

    # indices of the 3 nodes at the corner of the triangular
    node_indices = facet[1:]

    # coordinates of the corner nodes
    node_coords = coords[node_indices]

    # the centre point of the triangular facet
    face_centre = node_coords.mean(0)

    # indices of the element to which the surface facet belongs
    elem_index = facet[0]

    # indices of the 4 nodes at the corners of the tetrahedral element
    nodes_elem_index = elements[elem_index]

    # coordinates of the 4 corner nodes
    elem_nodes = coords[nodes_elem_index]

    # the centre point of the element
    elem_centre = elem_nodes.mean(0)

    # compute vectors which describe 3 sides of the triangular facet
    p0, p1, p2 = node_coords
    p1_0 = p1 - p0
    p2_0 = p2 - p0
    p2_1 = p2 - p1

    # find lengths of each side
    s0 = jnp.sqrt(((p1_0**2).sum()))
    s1 = jnp.sqrt(((p2_0**2).sum()))
    s2 = jnp.sqrt(((p2_1**2).sum()))

    # use Heron's formula to compute the area of the triangle
    s = (s0 + s1 + s2) / 2.
    area = jnp.sqrt(s*(s-s0)*(s-s1)*(s-s2))

    # the surface normal vector is orthognal to the plane in which the edge vectors lie
    Nt = jnp.cross(p2_0, p1_0)

    # rescale normal vector N to unit length
    N_mag = jnp.sqrt(((Nt**2).sum()))
    N = Nt / N_mag

    # ensure that N is pointing outwards
    v_in = elem_centre - face_centre
    N = jnp.where(jnp.dot(v_in, N) > 0, -N, N)

    return area, N

def compute_element_vol(element: jnp.ndarray, ref_coords: jnp.ndarray):
    """Compute volume of a tetrahedral element (reference configuration)"""

    # recentre coords to have origin equal to the first coord
    v10 = ref_coords[element[1]] - ref_coords[element[0]]
    v20 = ref_coords[element[2]] - ref_coords[element[0]]
    v30 = ref_coords[element[3]] - ref_coords[element[0]]

    # find tetrahedron volume using standard formula (see here for example: https://en.wikipedia.org/wiki/Tetrahedron#General_properties)
    v20_cross_v10 = jnp.cross(v20, v10)
    element_vol = jnp.abs(jnp.dot(v30, v20_cross_v10))/6.

    return element_vol

def compute_def_gradient_element(element: jnp.ndarray, ref_coords: jnp.ndarray, cur_coords:jnp.ndarray, Jtransform = lambda J: J):
    """Compute deformation gradient of a tetrahedral element given specified reference and current positions"""

    # extract reference and current coords of element
    node_0 = ref_coords[element]
    node_1 = cur_coords[element]

    # calculate the deformation gradient F as in Eq. (17) of the manuscript
    node_0_centered = (node_0[1:] - node_0[0]).T
    node_1_centered = (node_1[1:] - node_1[0]).T
    F = jnp.matmul(node_1_centered, jnp.linalg.inv(node_0_centered))

    # J is simply the determinant of F
    J = jnp.linalg.det(F)

    # apply transformation to J to prevent negative values or values close to zero (see Section 2.4.3 of the manuscript)
    J = Jtransform(J)

    return F, J


def compute_internal_work_element(F: jnp.ndarray, J: jnp.ndarray, elem_vol: jnp.ndarray, elem_fibre: jnp.ndarray, theta: jnp.ndarray, constitutive_law: Callable):
    """Compute internal work done in single element"""

    # internal work in single element computed as from first summand in Eq. (20) of the manuscript
    Psi_internal = constitutive_law(theta, F, J, elem_fibre) * elem_vol

    return Psi_internal


def compute_body_work_element(element: jnp.ndarray, elem_vol: jnp.ndarray, disp: jnp.ndarray, body_force: jnp.ndarray):
    """Compute body work done in single element"""

    # body work in single element computed as from second summand in Eq. (20) of the manuscript
    element_disp = disp[element].mean(0)
    body_force_work = (body_force * element_disp).sum() * elem_vol

    return body_force_work


def compute_surface_work_element(F: jnp.ndarray, J: jnp.ndarray, disp: jnp.ndarray, area_normals: jnp.ndarray, surface_force: jnp.ndarray):
    """Compute surface work done on single triangular surface facet"""

    # Nanson's formula (Eq. (18) of the manuscript)
    JNdA = -J*jnp.matmul(jnp.linalg.inv(F).T, area_normals)

    # $u_A^{fl}$ in the manuscript
    surface_disp = disp.mean(0)

    # surface work on single triangular facet computed as from third summand in Eq. (20) of the manuscript
    pressure_work = jnp.dot(surface_disp, JNdA) * surface_force

    return pressure_work


# vmap above functions which are defined per element/surface facet, to work over an entire mesh
compute_area_N        = jax.vmap(compute_area_N_facet, in_axes=[0, None, None])
compute_vol           = jax.vmap(compute_element_vol, [0, None])
compute_def_gradient  = jax.vmap(compute_def_gradient_element, in_axes = [0] + [None]*3)
compute_internal_work = jax.vmap(compute_internal_work_element, in_axes = [0]*4 + [None]*2)
compute_body_work     = jax.vmap(compute_body_work_element, in_axes = [0]*2 + [None]*2)
compute_surface_work  = jax.vmap(compute_surface_work_element, in_axes = [0]*4 + [None])


def total_potential_energy(displacement: jnp.ndarray, theta: jnp.ndarray, ref_geom_data, external_forces):
    """Compute total potential energy by evaluating Eq. (20) of the manuscript"""

    # current coords are simple reference coords + displacement
    cur_coords = ref_geom_data.ref_coords + displacement

    # compute deformation gradient for each element in mesh
    F, J = compute_def_gradient(ref_geom_data.elements, ref_geom_data.ref_coords, cur_coords, ref_geom_data.Jtransform)

    # compute internal work done given specified constitutive law
    Psi_internal = compute_internal_work(F, J, ref_geom_data.element_vols, ref_geom_data._fibre_field, theta, ref_geom_data.constitutive_law).sum()

    # initialise external work to be zero
    Psi_external = 0.

    # compute external work due to body forces
    if external_forces.body_force is not None:

        Psi_external += compute_body_work(ref_geom_data.elements,
                                          ref_geom_data.element_vols,
                                          displacement,
                                          external_forces.body_force).sum()

    # compute external work due to surface forces
    if external_forces.surface_force is not None:

        F_surface = F[ref_geom_data.tri_surface_element_indices]
        J_surface = J[ref_geom_data.tri_surface_element_indices]
        disp_surface = displacement[ref_geom_data.tri_surface_node_indices]

        surface_force = external_forces.surface_force

        Psi_external += compute_surface_work(F_surface,
                                             J_surface,
                                             disp_surface,
                                             ref_geom_data.tri_surface_area_normals,
                                             surface_force).sum()

    # return total potential energy
    return Psi_internal - Psi_external


utils_potential_energy.py
dodaltuin/soft-tissue-pignn
"""
File: utils_training.py
Author: David Dalton
Description: utility functionality for training and evaluating PI-GNN emulators
"""

import jax
import jax.numpy as jnp
from jax import jit, value_and_grad, vmap

import flax
import optax
from flax.training import train_state
from flax.metrics import tensorboard

# stops tensorflow from pre-allocating all of gpu memory
tensorboard.tf.config.experimental.set_visible_devices([], 'GPU')

from jax.config import config
config.update("jax_debug_nans", True)

import numpy as np

import pathlib
import pickle

from absl import logging
from functools import partial

import utils
import utils_data
import utils_evaluation as utils_eval
from utils_potential_energy import total_potential_energy


# MLP OPTIONS
MLP_WIDTH = 128
MLP_DEPTH = 2
MLP_FEATURES = [MLP_WIDTH] * MLP_DEPTH

# MeshGraphNets hyper-parameter values:
LOCAL_EMBED_DIM = 40

# Random seed for initialising network parameters
RNG_SEED = 2001

# optimiser for network parameters
OPTIMISATION_ALGORITHM = optax.adam

# whether to write prediction results arrays to disk after evaluation
SAVE_PREDICTIONS = True

# function to create subdir to save emulation results
create_savedir = partial(utils.create_savedir,
                                   local_embedding_dim=LOCAL_EMBED_DIM,
                                   mlp_width=MLP_WIDTH,
                                   mlp_depth=MLP_DEPTH,
                                   rng_seed=RNG_SEED)

# function to create configuration dictionary of emulator hyper-parameters
create_config_dict = partial(utils.create_config_dict,
                                       local_embed_dim=LOCAL_EMBED_DIM,
                                       mlp_features=MLP_FEATURES,
                                       rng_seed=RNG_SEED)

def compute_loss_pinn(params, theta_tuple, pred_fn, ref_geom_data, external_forces):
    """Compute total potential energy from emulator prediction"""

    theta_norm, theta = theta_tuple
    Upred = pred_fn(params, theta_norm)
    return total_potential_energy(Upred, theta, ref_geom_data, external_forces)

def train_step(params, opt_state, theta_tuple, optimiser, loss_fn):
    """Train emulator for one theta input point """

    partial_loss_fn = partial(loss_fn, theta_tuple=theta_tuple)
    grad_fn = jax.value_and_grad(partial_loss_fn)
    loss, grads = grad_fn(params)
    updates, opt_state = optimiser.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

def predict_dataset(data_loader, pred_fn):
    """Make predictions for entire dataset"""

    Upred = np.zeros_like(data_loader._displacement)
    for graph_idx in data_loader._epoch_indices:
        input_output_tuple = data_loader.get_graph(graph_idx)
        theta_norm, _, _, _ = input_output_tuple
        Upred[graph_idx] = np.array(pred_fn(theta_norm))
    return Upred

class PhysicsLearner:
    """Class for training PI-GNN emulator and saving learned parameters"""

    def __init__(self, pred_fn, train_dg, params, lr, optim_algorithm, ref_geom_data, external_forces, logging, results_save_dir = None, summary_writer=None):

        self.train_dg = train_dg
        self.params = params
        self.lr = lr
        self.optim_algorithm = optim_algorithm
        self.init_optimiser()
        self.opt_state = self.optimiser.init(params)
        self.n_epochs_trained = 0
        self.offset_idx = 0
        self.logging = logging
        self.summary_writer = summary_writer
        self.results_save_dir = results_save_dir
        self.min_train_loss = 1e7

        # intitialise loss as function of displacement and theta
        self.train_loss_fn = partial(compute_loss_pinn,
                                     pred_fn=pred_fn,
                                     ref_geom_data=ref_geom_data,
                                     external_forces=external_forces)

        # jit the training step function for faster execution
        self.train_step = jit(partial(train_step,
                                      optimiser = self.optimiser,
                                      loss_fn = self.train_loss_fn))

    def train_epoch(self, random_sampling=False):
        """Train network for one epoch"""

        if random_sampling:
            self.train_dg.resample_input_points()
        else:
            self.train_dg.shuffle_epoch_indices()

        loss = 0.
        for graph_idx in self.train_dg.epoch_indices:
            theta_tuple_idx = self.train_dg.get_data(graph_idx)
            self.params, self.opt_state, loss_idx = self.train_step(self.params, self.opt_state, theta_tuple_idx)
            loss += loss_idx
        # train loss for epoch is mean total potential energy
        self.train_loss = loss / self.train_dg.epoch_size

    def fit_pinn(self, n_epochs: int, save_params = False, random_sampling=False):
        """Train network for 'n_epochs' epochs"""

        self.logging.info(f'Beginning training for {n_epochs} epochs')
        for epoch_idx in range(n_epochs):

            # train network for one epoch
            self.train_epoch(random_sampling)

            # keep track of number of training epochs that have been completed
            self.n_epochs_trained += 1

            # offset epoch_idx to account for any previous calls to self.fit_pinn
            epoch_idx_total = self.offset_idx + epoch_idx

            # save trained network parameters based on validation set prediction error
            if save_params:
                self.save_trained_params(self.train_loss)

            # write loss values to tensorboard summary_writer
            if self.summary_writer is not None:
                self.summary_writer.scalar('train_loss', self.train_loss, epoch_idx_total)
                self.summary_writer.scalar('learning_rate', self.opt_state.hyperparams["learning_rate"], epoch_idx_total)

            if (epoch_idx % 250 == 0) or (epoch_idx < 150):
                self.logging.info(f'({epoch_idx_total}): train_loss={self.train_loss:.5f}, lr={self.opt_state.hyperparams["learning_rate"]:.1e}')

        # keep track of number of training epochs that have been performed for reference if training is restarted later
        self.offset_idx = self.n_epochs_trained

    def init_optimiser(self):
        """Initialise the optimiser used for training"""

        self.optimiser = optax.inject_hyperparams(self.optim_algorithm)(learning_rate=self.lr)

    def update_learning_rate(self, new_lr):
        """Update learning rate used for training"""

        self.lr = new_lr
        self.init_optimiser()
        self.opt_state = self.optimiser.init(self.params)
        self.train_step = jit(partial(train_step,
                                      optimiser = self.optimiser,
                                      loss_fn = self.train_loss_fn))


    def save_trained_params(self, epoch_loss):
        """Save network parameters if current loss exceeds minimum loss"""

        if epoch_loss < self.min_train_loss:
            self.min_train_loss = epoch_loss
            with pathlib.Path(self.results_save_dir, f'trainedNetworkParams.pkl').open('wb') as fp:
                pickle.dump(self.params, fp)


def train(data_path: str, K: int, n_epochs: int, lr: float, trained_params_dir: str, dir_label: str):
    """Train a PI-GNN emulator and save results"""

    # create directory to store the trained parameters of the network
    results_save_dir = create_savedir(data_path, K, n_epochs, lr, dir_label)

    logging.get_absl_handler().use_absl_log_file('training', f'{results_save_dir}/logFiles')
    logging.set_stderrthreshold(logging.DEBUG)

    logging.info(f'Beginning Training')
    logging.info(f'Training epochs: {n_epochs}')
    logging.info(f'Learning rate: {lr}')
    logging.info(f'Data path: {data_path}')
    logging.info(f'Message passing steps (K): {K}')
    logging.info(f'Trained Params Dir: {trained_params_dir}\n')
    logging.info(f'Results save directory: {results_save_dir}\n')

    # load reference geometry data
    ref_geom = utils_data.ReferenceGeometry(data_path, logging)

    # store external force data (body forces and/or surface forces)
    external_forces = utils_data.ExternalForces(data_path)

    # initialise generator of training input points
    train_dg = utils_data.DataGenerator(data_path)

    # create configuration dictionary of hyperparameters of the GNN emulator
    config_dict = create_config_dict(K, n_epochs, lr, ref_geom._output_dim)

    # write configuration dictionary to a text file in results_save_dir
    with open(f'{results_save_dir}/config_dict.txt', 'w') as f: print(config_dict, file=f)

    # create tensorboard SummaryWriter to monitor training
    summary_writer = tensorboard.SummaryWriter(results_save_dir)
    summary_writer.hparams(config_dict)

    # initialise GNN emulator, trainable parameters and emulator forward predict function
    pred_fn, params, emulator = utils.initialise_emulator(config_dict, train_dg, results_save_dir, ref_geom, fix_message_layers=False)

    # zero out the weights in the last layer of the decoder FCNNs
    params = utils.gen_zero_params_gnn(emulator, params)

    # initialise learning for physics-informed training
    learner = PhysicsLearner(pred_fn, train_dg, params, lr, OPTIMISATION_ALGORITHM, ref_geom, external_forces, logging, results_save_dir, summary_writer)

    # train first half at learning rate lr
    n_epochs_start = max(n_epochs//2, 1)
    learner.fit_pinn(n_epochs_start, save_params=True, random_sampling=False)

    # finish training at learning rate lr/10
    learner.update_learning_rate(lr/10.)
    n_epochs_end = max(n_epochs-n_epochs_start, 1)
    learner.fit_pinn(n_epochs_end, save_params=False, random_sampling=True) # randomly sample material parameters at each step

    # save trained network params
    with pathlib.Path(learner.results_save_dir, f'trainedNetworkParams.pkl').open('wb') as fp:
         pickle.dump(learner.params, fp)


def evaluate(data_path: str, K: int, n_epochs: int, lr: float, trained_params_dir: str, dir_label: str):
    """Evaluate performance of a trained a PI-GNN emulator on simulation data"""

    # directory where results are saved
    results_save_dir = create_savedir(data_path, K, n_epochs, lr, dir_label)

    logging.get_absl_handler().use_absl_log_file('evaluation', f'{results_save_dir}/logFiles')
    logging.set_stderrthreshold(logging.DEBUG)

    logging.info('Beginning Evaluation')
    logging.info(f'Data path: {data_path}')
    logging.info(f'Message passing steps (K): {K}')
    logging.info(f'Training epochs: {n_epochs}')
    logging.info(f'Learning rate: {lr}')
    logging.info(f'Trained Params Dir: {trained_params_dir}\n')
    logging.info(f'Results save directory: {results_save_dir}\n')

    # load reference geometry data
    ref_geom = utils_data.ReferenceGeometry(data_path)

    # store external force data (body forces and/or surface forces)
    external_forces = utils_data.ExternalForces(data_path)

    # load test simulation data
    test_data = utils_data.DataLoader(data_path, 'test')
    logging.info(f'Number of test data points: {test_data._data_size}')

    # create dictionary of hyperparameters of the GNN emulator
    config_dict = create_config_dict(K, n_epochs, lr, ref_geom._output_dim)

    # if trained_params_dir is not set, parameters are read from results_save_dir
    if trained_params_dir == "None": trained_params_dir = results_save_dir

    # initialise GNN emulator and read trained network parameters
    pred_fn, trained_params, emulator = utils.initialise_emulator(config_dict, test_data, results_save_dir, ref_geom, True, trained_params_dir)

    # vmap to allow total potential energy to be computed for all simulations similtaneously
    pe_vmap = jax.vmap(partial(total_potential_energy, ref_geom_data=ref_geom, external_forces=external_forces))

    # hardcode trained parameters into prediction function and jit for faster execution
    pred_fn_jit = jit(lambda theta_norm: pred_fn(trained_params, theta_norm))

    logging.info('Predicting on test data set using trained emulator')
    Upred = predict_dataset(test_data, pred_fn_jit)

    # calculate total potential energy for test data using emulator prediction
    PEpred = pe_vmap(Upred,test_data._theta)

    # retrieve the corresponding values obtained using the finite-element simulator
    Utrue = test_data._displacement
    PEtrue = test_data._pe_values

    # find error between total potential energy calculated in JAX versus results returned by FEniCS (PEtrue)
    PEtrue_calc = pe_vmap(Utrue,test_data._theta)
    calc_errors = 100.*jnp.abs((PEtrue - PEtrue_calc)/PEtrue)
    logging.info(f'Max PE calc error against FEniCS: {calc_errors.max():.4f}% (sim {calc_errors.argmax()})\n')

    # compute deformation gradient and potential energy from true displacements
    Ftrue, Jtrue, I1true = utils_eval.compute_F_J_I1_vmap(Utrue, ref_geom)

    # compute deformation gradient and potential energy from predicted displacements
    Fpred, Jpred, I1pred = utils_eval.compute_F_J_I1_vmap(Upred, ref_geom)

    # collect true/predicted arrays into tuples
    true_arrs = Utrue, PEtrue, Ftrue, Jtrue, I1true
    pred_arrs = Upred, PEpred, Fpred, Jpred, I1pred

    # print prediction error statistics to console
    utils_eval.print_error_statistics(true_arrs, pred_arrs, results_save_dir, logging)

    if SAVE_PREDICTIONS:
        logging.info('Saving Results')
        np.save(f'{results_save_dir}/predDisplacement.npy', Upred)
        np.save(f'{results_save_dir}/trueDisplacement.npy', Utrue)

    # make Paraview (.vtk) and Augmented Reality (.ply) visualisations of results
    try:
        import paraview.simple
    except:
        logging.warning('3D visualisations cannot be generated, as there was en error in importing paraview. Try typing "from paraview.simple import *" in a Python shell and follow the error printed to the screen.')
        return

    from utils_visualisation import make_3D_visualisations as make_visuals
    logging.info('Generating 3D visualisation files')
    make_visuals(Utrue, Upred, ref_geom.ref_coords, data_path, results_save_dir, logging)


# store trainig and evaluation functions in dictionary for access in "main.py"
run_fns_dict = {'train': [train],
                'evaluate': [evaluate],
                'train_and_evaluate': [train, evaluate]}

utils_training.py
dodaltuin/soft-tissue-pignn
"""
File: utils_visualisation.py
Author: David Dalton
Description: Generate .vtk and .ply files for visualisation of emulation results
"""

from paraview.simple import *
import numpy as np
import shutil
import os
import h5py
from absl import logging

h5_filename   = "sim_output.h5"
xdmf_filename = "sim_output.xdmf"

def find_quantile_indices(loss_array):
    """Find indices of the median and worst emulation errors"""

    # median (worst) error is 50th (100th) percentile
    loss_quantiles_labels = ['Median', 'Worst']
    loss_quantiles = [50, 100]

    # calculate the values which give rise to median/worst values
    loss_quantile_values = np.percentile(loss_array, loss_quantiles)

    # write results to dictionary
    index_dict = {}

    # loop over median and worst emulation results
    for i in range(len(loss_quantiles)):

        label_i = loss_quantiles_labels[i]
        loss_quantile_i = loss_quantile_values[i]

        # find the simulation which is closest to the quantile value calculated above
        index_i = np.argmin(np.abs(loss_array - loss_quantile_i))

        # write results to dictionary
        index_dict[label_i] = index_i

    return index_dict

def generate_savedir(label, index, results_save_dir):
    """Generate directory where visualisation files will be saved"""
    return f'{results_save_dir}resultsVisualisations/{label}_index{index}'

def copy_across_files(index_dict, true_disp, pred_disp, ref_coords, results_save_dir, data_save_dir):
    """Copy FEniCS template files and .npy results files to savedir"""

    # loop over median and worst emulation results
    for label, index in index_dict.items():

        # create directory to visualistion files
        savedir = generate_savedir(label, index, results_save_dir)
        if not os.path.isdir(savedir): os.makedirs(savedir)

        # save np arrays of results for this simulation index
        np.save(f'{savedir}/ref_geom.npy', ref_coords)
        np.save(f'{savedir}/true_disp.npy', true_disp[index])
        np.save(f'{savedir}/pred_disp.npy', pred_disp[index])
        np.save(f'{savedir}/losses.npy', true_disp[index] - pred_disp[index])

        #assert os.path.isfile(f'{data_save_dir}/sim_output.h5') and os.path.isfile(f'{data_save_dir}/sim_output.xdmf'), f"The files sim_output.h5 and sim_output.xdmf must exist in {data_save_dir} so that 3D visualisations can be made"

        # copy across FEniCS simulation output files for temporary writing of above information
        fenics_files_exist = os.path.isfile(f'{data_save_dir}/sim_output.h5') and os.path.isfile(f'{data_save_dir}/sim_output.xdmf')
        if not fenics_files_exist:
            logging.info(f"Warning - no 3D visualisations made because the template files sim_output.h5 and sim_output.xdmf do not exist in {data_save_dir}. For details on how these files can be generated, see data/DATA_FORMAT_REQUIREMENTS.md")
            exit()

        shutil.copy(f'{data_save_dir}/sim_output.h5',   f'{savedir}/{h5_filename}')
        shutil.copy(f'{data_save_dir}/sim_output.xdmf', f'{savedir}/{xdmf_filename}')


def generate_visualisation_files(index_dict, root_save_dir):
    """Generate .vtk and .ply visualisation files"""

    # loop over median and worst emulation results
    for label, index in index_dict.items():

        # create directory to visualistion files
        savedir = generate_savedir(label, index, root_save_dir)

        # find names of the npy files in savedir (without .npy extension)
        npy_files = [f[:-4] for f in os.listdir(savedir) if f.endswith('.npy')]


        for file_i in npy_files:

            # load the data from the .npy file generated in "copy_across_files"
            np_arr = np.load(f'{savedir}/{file_i}.npy')

            # create zero displacement for reference geometry for consistent plotting
            if file_i=="ref_geom": np_arr*=0.

            # write np array values to .h5 file
            f1 = h5py.File(f'{savedir}/{h5_filename}', "r+")
            f1['VisualisationVector']['0'].write_direct(np_arr)
            f1.close()

            # read .xdmf mesh file from fenics
            xdmf_geometry = XDMFReader(registrationName=xdmf_filename, FileNames=[f'{savedir}/{xdmf_filename}'])
            point_labels  = xdmf_geometry.PointArrayStatus[0]

            # extract surface from input geometry
            surface_model = ExtractSurface(registrationName='surface_model', Input=xdmf_geometry)

            # create reference geometry variable by extracting surface normals
            reference_configuration = GenerateSurfaceNormals(registrationName='reference_configuration', Input=surface_model)

            # update reference configuration by warping with given vectors
            updated_configuration = WarpByVector(registrationName='updated_configuration', Input=reference_configuration)
            updated_configuration.Vectors = ['POINTS', point_labels]

            # get colour map from the displacement vector magnitude
            colour_map = GetColorTransferFunction(point_labels)

            # save vtk files for paraview visualisation
            SaveData(f'{savedir}/{file_i}.vtk', proxy=updated_configuration, PointDataArrays=['Normals', point_labels])

            # save ply files for Augmented Reality visualistions
            SaveData(f'{savedir}/{file_i}.ply', proxy=updated_configuration, PointDataArrays=['Normals', point_labels],
                EnableColoring=1,
                ColorArrayName=['POINTS', point_labels],
                LookupTable=colour_map)


def make_3D_visualisations(Utrue, Upred, ref_coords, data_path, results_save_dir, logging):
    """Wrapper function which calls each of the above functions in turn"""

    pred_errors = (((Utrue - Upred)**2).sum(-1)**.5).mean(-1)

    index_dict = find_quantile_indices(pred_errors)

    logging.info(f'Index of simulation with median (worst) prediction errors: {index_dict["Median"]} ({index_dict["Worst"]})')

    copy_across_files(index_dict, Utrue, Upred, ref_coords, results_save_dir,f'data/{data_path}/geometryData')

    generate_visualisation_files(index_dict, results_save_dir)


utils_visualisation.py
dodaltuin/soft-tissue-pignn