implement hash grid

README.md

@@ -106,27 +106,33 @@ $ tree . -L 1 --filesfirst
 ```
 
 ```bash
-$ tree src/ -L 1
-├── autodiff.jl        # AD wrapper for 1-4th order derivatives
-├── evolve.jl          # Logic for dynamics evaluation
-├── layers.jl          # Architecture definitions (e.g., auto-decode, C-ROM, SIREN, etc.)
-├── metrics.jl         # Loss functions
-├── neuralgridmodel.jl # Grid-dependent neural space discretization (e.g., CAE-ROM, POD-ROM)
-├── neuralmodel.jl     # Neural field spatial discretization (e.g., C-ROM, SNF-ROM)
-├── NeuralROMs.jl      # Main file: declares Julia module and imports relevant packages
-├── nonlinleastsq.jl   # Nonlinear least square solve for LSPG and for initializing auto-decode.
-├── operator.jl        # Operator kernel definitions
-├── optimisers.jl      # Modified weight decay optimisers based on Optimisers.jl
-├── problems.jl        # PDE problem definitions du/dt = f(x, u, t, u', ...)
-├── timeintegrator.jl  # Time integrator object definition
-├── train.jl           # Training script
-├── transform.jl       # Fourier transform for Fourier Neural Operator
-├── utils.jl           # Miscalleneous utility functions
-└── vis.jl             # 1D/2D visualization functions
+$ tree src/ -L 2 --filesfirst
+.
+├── autodiff.jl            # AD wrapper for 1-4th order derivatives
+├── metrics.jl             # Loss functions
+├── neuralgridmodel.jl     # Grid-dependent neural space discretization (e.g., CAE-ROM, POD-ROM)
+├── neuralmodel.jl         # Neural field spatial discretization (e.g., C-ROM, SNF-ROM)
+├── NeuralROMs.jl          # Main file: declares Julia module and imports relevant packages
+├── nonlinleastsq.jl       # Nonlinear least square solve for LSPG and for initializing auto-decode.
+├── optimisers.jl          # Modified weight decay optimisers based on Optimisers.jl
+├── pdeproblems.jl         # PDE problem definitions du/dt = f(x, u, t, u', ...)
+├── train.jl               # Training loop
+├── utils.jl               # Miscalleneous utility functions
+├── vis.jl                 # 1D/2D visualizations
+├── dynamics               #
+│   ├── evolve.jl          # Logic for dynamics evaluation
+│   └── timeintegrator.jl  # Time integrator object definition
+├── layers                 #
+│   ├── basic.jl           # Basic layer definitions (e.g., PermuteLayer, HyperNet)
+│   ├── encoder_decoder.jl # Encoder-decoder network definitions (auto-decode, CAE, C-ROM, SNF-ROM)
+│   └── sdf.jl             # Layers for 3D shape encoding
+└── operator               #
+    ├── oplayers.jl        # Fourier neural operator kernel definitions
+    └── transform.jl       # Spectral transforms for FNO
 ```
 
 ```bash
-$ tree examples/ -L 1 --filesfirst
+$ tree experiments_SNFROM/ -L 1 --filesfirst
 experiments_SNFROM/
 ├── autodecode.jl     # Autodecode-ROM training and inference
 ├── cases.jl          # Experiment setup

experiments_SDF/Project.toml

@@ -2,5 +2,7 @@
 Bonito = "824d6782-a2ef-11e9-3a09-e5662e0c26f8"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
 MeshCat = "283c5d60-a78f-5afe-a0af-af636b173e11"
+Meshes = "eacbb407-ea5a-433e-ab97-5258b1ca43fa"
 Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73"
 WGLMakie = "276b4fcb-3e11-5398-bf8b-a0c2d153d008"
+WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"

experiments_SDF/SDF.jl

@@ -14,9 +14,9 @@ let
     !(sdfpath in LOAD_PATH) && push!(LOAD_PATH, sdfpath)
 end
 
-using MeshCat
-using GeometryBasics: Mesh, HyperRectangle, Vec, SVector
-using Meshing: isosurface, MarchingCubes, MarchingTetrahedra, NaiveSurfaceNets
+# using MeshCat
+# using GeometryBasics: Mesh, HyperRectangle, Vec, SVector
+# using Meshing: isosurface, MarchingCubes, MarchingTetrahedra, NaiveSurfaceNets
 
 using WGLMakie, Bonito
 
@@ -178,6 +178,7 @@ function postprocess_SDF(
     model = jldopen(modelfile)
     NN, p, st = model["model"]
     md = model["metadata"]
+    δ = md.δ
 
     display(NN)
     @show md
@@ -200,6 +201,7 @@ function postprocess_SDF(
         x = vcat(xx', zz', yy')
         @time u = eval_model((NN, p, st), x; device) #, batchsize = 100)
         @time u = reshape(u, samples...)
+        u = Float16.(u)
 
         file = h5open(sdffile, "w")
         write(file, "u", u)
@@ -208,10 +210,12 @@ function postprocess_SDF(
         u
     end
 
+    # # MeshCat workflow
+
     # # RUN.JL
     # isdefined(Main, :vis) && MeshCat.close_server!(vis.core)
     # vis = Visualizer()
-    # postprocess_SDF(casename, modelfile; vis, device)
+    # vis = postprocess_SDF(casename, modelfile; vis, device)
     # open(vis; start_browser = false)
 
     # # HERE
@@ -220,17 +224,17 @@ function postprocess_SDF(
     # setobject!(vis, mesh)
     # return vis
 
-    # u = Float16.(u)
+    # # WGLMakie workflow
+
     # plotsdf() = volume(u;
     #     algorithm = :iso,
-    #     isovalue = 0.00f0,
-    #     isorange = 0.002f0,
+    #     isovalue = zero(Float16),
+    #     isorange = Float16(0.10δ),
     #     nan_color = :transparent,
     #     colormap = [:transparent, :gray]
     # )
 
     u = map(x -> x ≤ 0 ? 1 : 0, u) |> BitArray
-
     plotsdf() = volume(u;
         absorption = 50,
         algorithm = :absorption,

experiments_SDF/run.jl

@@ -15,11 +15,11 @@ Random.seed!(rng, 199)
 # casename = "Burger.npz"
 # modeldir  = joinpath(@__DIR__, "dump3")
 
-# casename = "HumanSkull.npz"
-# modeldir  = joinpath(@__DIR__, "dump4")
+casename = "HumanSkull.npz"
+modeldir  = joinpath(@__DIR__, "dump4")
 
-casename = "Cybertruck.npz"
-modeldir  = joinpath(@__DIR__, "dump5")
+# casename = "Cybertruck.npz"
+# modeldir  = joinpath(@__DIR__, "dump5")
 
 modelfile = joinpath(modeldir, "model_08.jld2")
 device = Lux.gpu_device()
@@ -31,6 +31,6 @@ h, w = 5, 512
 # isdir(modeldir) && rm(modeldir, recursive = true)
 # model, ST, md = train_SDF(casename, modeldir, h, w, E; rng, device)
 
-isdefined(Main, :server) && close(server)
-server = postprocess_SDF(modelfile; device)
+# isdefined(Main, :server) && close(server)
+# server = postprocess_SDF(modelfile; device)
 #======================================================#

experiments_SDF/utils.jl

@@ -1,32 +1,4 @@
 #
-#===========================================================#
-# Multiresolution has encoding
-#===========================================================#
-
-"""
-https://github.com/cheind/pure-torch-ngp/blob/develop/torchngp/modules/encoding.py
-"""
-function index_hash(
-    Ix::AbstractVector{<:Integer}, # [K]
-    Iy::AbstractVector{<:Integer},
-    Iz::AbstractVector{<:Integer},
-    nEncodings::Integer,
-)
-
-    primes = [1, 2654435761, 805459861]
-
-    @. xor(Ix * primes[1], Iy * primes[2], Iz * primes[3]) % nEncodings
-end
-
-function index_dense(
-    Ix::AbstractVector{<:Integer},
-    Iy::AbstractVector{<:Integer},
-    Iz::AbstractVector{<:Integer},
-    shape::NTuple{3, <:Integer},
-)
-    Iz * (shape[1] * shape[2]) + Iy * shape[1] + Ix
-end
-
 #===========================================================#
 function make_optimizer(
     E::Integer,

experiments_SDF/view.jl

@@ -0,0 +1,64 @@
+#
+using NeuralROMs, NPZ
+
+let
+    pkgpath = dirname(dirname(pathof(NeuralROMs)))
+    sdfpath = joinpath(pkgpath, "experiments_SDF")
+    !(sdfpath in LOAD_PATH) && push!(LOAD_PATH, sdfpath)
+end
+
+#======================================================#
+
+basedir  = joinpath(pkgdir(NeuralROMs), "experiments_SDF", "dataset_netfabb")
+datafile = joinpath(basedir, "56250_3b6024e3_54.npz")
+
+data = npzread(datafile)
+verts = data["verts"]' # [3, V]
+elems = data["elems"]' # [8, F] (hex mesh)
+displ = data["disp"]'  # [3, V]
+
+import Meshes
+using GeometryBasics
+using GeometryBasics: Mesh
+
+verts = [Point3f(verts[:, i]...) for i in axes(verts, 2)]
+
+elems = vcat(
+    elems[1, :]',
+    elems[2, :]',
+    elems[3, :]',
+    elems[4, :]',
+    #
+    elems[7, :]',
+    elems[8, :]',
+    elems[5, :]',
+    elems[6, :]',
+)
+
+faces1 = vcat(elems[1, :]', elems[2, :]', elems[3, :]', elems[4, :]') # bot
+faces2 = vcat(elems[1, :]', elems[2, :]', elems[6, :]', elems[5, :]')
+faces3 = vcat(elems[2, :]', elems[3, :]', elems[7, :]', elems[6, :]')
+faces4 = vcat(elems[3, :]', elems[4, :]', elems[8, :]', elems[7, :]')
+faces5 = vcat(elems[4, :]', elems[1, :]', elems[5, :]', elems[8, :]')
+faces6 = vcat(elems[5, :]', elems[6, :]', elems[7, :]', elems[8, :]') # top
+
+faces = hcat(faces1, faces2, faces3, faces4, faces5, faces6)
+mesh  = Mesh(verts, vec(faces), QuadFace)
+
+using MeshCat
+isdefined(Main, :vis) && MeshCat.close_server!(vis.core)
+vis = Visualizer()
+setobject!(vis, mesh)
+open(vis; start_browser = false)
+
+# using WriteVTK
+# cell_type = VTKCellTypes.VTK_HEXAHEDRON
+# cells = [MeshCell(cell_type, elems[:, i]) for i in axes(elems, 2)]
+#
+# outfile = "outfile"
+# vtk_out = vtk_grid(outfile, verts, cells)
+# # vtk_out["displacement", VTKPointData()] =  displ
+# close(vtk_out)
+
+#======================================================#
+nothing

src/NeuralROMs.jl

@@ -79,22 +79,24 @@ export
     forwarddiff_deriv1, forwarddiff_deriv2, forwarddiff_deriv4, forwarddiff_jacobian,
     finitediff_deriv1, finitediff_deriv2, finitediff_deriv4, finitediff_jacobian
 
-include("layers.jl")
+include("layers/basic.jl")
+export Atten, Diag, PermutedBatchNorm, SplitRows, HyperNet
+
+include("layers/encoder_decoder.jl")
 export
-    Atten, Diag, PermutedBatchNorm, SplitRows, ImplicitEncoderDecoder,
-    HyperNet,
+    ImplicitEncoderDecoder,
     AutoDecoder, get_autodecoder,
     FlatDecoder, get_flatdecoder, freeze_decoder,
     HyperDecoder, get_hyperdecoder
 
-include("transform.jl")
-export FourierTransform, CosineTransform
+include("layers/sdf.jl")
+export ClampVanilla, ClampTanh, ClampSigmoid, ClampSoftsign
 
-include("operator.jl")
+include("operator/oplayers.jl")
 export OpKernel, OpConv, OpKernelBilinear, OpConvBilinear, linear_nonlinear
 
-include("sdf.jl")
-export ClampVanilla, ClampTanh, ClampSigmoid, ClampSoftsign
+include("operator/transform.jl")
+export FourierTransform, CosineTransform
 
 include("optimisers.jl")
 export DecoderWeightDecay, IdxWeightDecay
@@ -109,7 +111,7 @@ export
 
 include("neuralgridmodel.jl")
 
-include("problems.jl")
+include("pdeproblems.jl")
 export dudtRHS
 export
     Advection1D, Advection2D,
@@ -121,10 +123,10 @@ export
 include("nonlinleastsq.jl")
 export nonlinleastsq
 
-include("timeintegrator.jl")
+include("dynamics/timeintegrator.jl")
 export TimeIntegrator, perform_timestep!, evolve_integrator!, evolve_model
 
-include("evolve.jl")
+include("dynamics/evolve.jl")
 export
     # timestepper types
     EulerForward, EulerBackward, RK2, RK4,