jax_fem.utils module

jax_fem.utils module#

jax_fem.utils.save_sol(fe, sol, sol_file, cell_infos=None, point_infos=None)[source]#

Save finite element solution and associated data to VTK file.

Parameters:

  • fe (FiniteElement) – Finite element object.

  • sol (JaxArray) – Solution vector to save (vertex-based). Shape is (num_total_nodes, vec).

  • sol_file (str) – Output file path.

  • cell_infos (list) –

    Additional cell data as [(name1, data1), (name2, data2)]. Each data array must have shape (num_cells,…). For example,

    cell_infos = [ ('p', p_cell_data)]

  • point_infos (list) –

    Additional point data as [(name1, data1), (name2, data2)]. Each data array must have shape (num_total_nodes,…).

    For example,

    point_infos = [('T', T_point_data)]

jax_fem.utils.read_abaqus_and_write_vtk(abaqus_file, vtk_file)[source]#

Used for a quick inspection. Paraview can’t open .inp file so we convert it to .vtu

Parameters:

  • abaqus_file (str) – Input Abaqus .inp file path

  • vtk_file (str) – Output VTK file path (.vtu or .vtk)

jax_fem.utils.json_parse(json_filepath)[source]#

Parse JSON configuration file and print formatted contents.

Parameters:

json_filepath (str) – Path to JSON configuration file

Returns:

Parsed JSON data as dictionary

Return type:

dict

jax_fem.utils.make_video(data_dir)[source]#

Generate MP4 video from PNG sequence using ffmpeg.

  • The command -pix_fmt yuv420p is to ensure preview of video on Mac OS is enabled (see ref1).

  • The command -vf "pad=ceil(iw/2)*2:ceil(ih/2)*2" is to solve the “not-divisible-by-2” problem (see ref2.).

  • The command -y means always overwrite.

Parameters:

data_dir (str) – Directory containing PNG frames in {data_dir}/png/tmp/

Notes

Output saved as {data_dir}/mp4/test.mp4. Requires ffmpeg.

jax_fem.utils.timeit(func)[source]#

Decorator for printing the timing results of a function.

Parameters:

func (callable) – Function to be timed.

Returns:

Wrapped function with timing logic.

Return type:

callable

jax_fem.utils.walltime(txt_dir=None, filename=None)[source]#

Wrapper for writing timing results to a file

Parameters:

  • txt_dir (str) – Directory to save timing data.

  • filename (str) – Base filename (default: ‘walltime_{platform}.txt’).

Returns:

Decorator function.

Return type:

callable

Contents