init ML solver

Project.toml

@@ -18,20 +18,27 @@ FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 PaddedViews = "5432bcbf-9aad-5242-b902-cca2824c8663"
 ParallelOperations = "09c1cff2-b94d-4c31-85a1-721512e21c63"
 PhysicalFDM = "4fd04e27-7767-4ec9-b70f-09a3b509a447"
+PhysicalFFT = "d41fe92e-7e57-414e-90fb-fae81d34872f"
 PhysicalMeshes = "97d9904f-034f-4fb7-aeaa-03a173434233"
 PhysicalParticles = "78724250-4eb2-11e9-0f34-13a7aab864c5"
 PhysicalTrees = "dc5d4990-50a5-4af2-bdaa-6f78dc9b9fb2"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

src/AstroNbodySim.jl

@@ -17,6 +17,8 @@ using DataFrames
 using DocStringExtensions
 using Base.Iterators
 using Measurements
+using JLD2
+using Random
 
 # user-friendly
 using Unitful, UnitfulAstro
@@ -45,18 +47,22 @@ using GLMakie
 using LoopVectorization
 
 # Machine Learning
-# using Knet
+using Lux, LuxCUDA
+using MLUtils
+using Optimisers
 using Images
 using Zygote
 
 # JuliaAstroSim
 using ParallelOperations
+@reexport using AstroSimBase
 @reexport using PhysicalParticles
 @reexport using PhysicalMeshes
 @reexport using PhysicalTrees
 @reexport using AstroIO
 @reexport using PhysicalFDM
-@reexport using AstroSimBase
+@reexport using PhysicalFFT
+using PhysicalParticles.NumericalIntegration
 
 # GPU
 using CUDA
@@ -137,8 +143,8 @@ export
     delta_mat2,
     delta_mat3,
     laplace_conv_op, laplace_conv,
-    fft_poisson, fft_poisson!,
     fdm_poisson,
+    compute_acc,
 
     # ML
     # train_cnn_poisson2d,
@@ -250,13 +256,11 @@ include("directsumgpu/timestep.jl")
 
 include("PM/gravity.jl")
 include("PM/fft.jl")
-include("PM/cnn.jl")
 include("PM/output.jl")
 include("PM/timestep.jl")
 
-# include("ML/cnn-poisson/models.jl")
-# include("ML/cnn-poisson/dataset.jl")
-# include("ML/cnn-poisson/train.jl")
+include("ML/training.jl")
+include("ML/gravity.jl")
 
 include("restart.jl")
 include("run.jl")