Transition to Pythoncall from PyCall (#135)

* Update Project with new dependencies

* remove FiniteDifferences from dependencies

* Example of APWP fit based on Jupp1987

* add complex-step method

* Double differentiation working with complex-step method

* Initial condition u0 fitting implemented

* Projected gradient descent working for u0. Some more tests

* Testing activation functions with complex-step

* predict function, return multiple losses

* Co-authored-by: Jordi Bolibar <jordi.bolibar@gmail.com>

* Multiple shooting working once sensealg specified

* Double rotation example with small changes in src

* feat: update to support Lux 1.0

* Example with double rotation working with non-updated Lux

* Integration test of inversion

* Added Random as test dependency

* Fix Lux version

Co-authored-by: Avik Pal <avikpal@mit.edu>

* Reorganization of loss function in different module

* Implementation of cubic splines as Jupp 1987

* Curl examples and reweighted loss experiment

* minimal changes

* Bump version

* Added dependencies to compat

* Remove OptimizationPolyalgorithms from dependencies

* Updated examples

* Transition to PythonCall

* Update CI.yml - Simplify workflow without PyCall

* add PythonCall to dependencies

* Update Project.toml - bump version

---------

Co-authored-by: Avik Pal <avikpal@mit.edu>

.github/workflows/CI.yml

@@ -3,7 +3,6 @@ on:
   push:
     branches:
       - main
-      - up-lux
     tags: ['*']
   pull_request:
     branches:
@@ -40,26 +39,6 @@ jobs:
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python }}  
-      - name: Create environment with micromamba 🐍🖤
-        uses: mamba-org/setup-micromamba@v1
-        with: 
-          micromamba-version: '1.3.1-0'
-          environment-file: ./environment.yml
-          environment-name: SphereUDE
-          init-shell: bash
-          cache-environment: true
-      - name: Test creation of environment with micromamba 🔧🐍🖤
-        run: | 
-            which python
-            conda env export
-        shell: bash -el {0}
-      - name: Update certifi
-        run: | 
-            pip install --upgrade certifi
-        shell: bash -el {0}
-      - name: Set ENV Variables for PyCall.jl 🐍 📞
-        run: export PYTHON=/home/runner/micromamba/envs/SphereUDE/bin/python
-        shell: bash -el {0}
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
@@ -69,8 +48,6 @@ jobs:
           cache-registries: "true"
           cache-compiled: "true"
       - uses: julia-actions/julia-buildpkg@v1
-        env:
-          PYTHON : /home/runner/micromamba/envs/SphereUDE/bin/python
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v3

CondaPkg.toml

@@ -0,0 +1,13 @@
+[deps]
+future = ""
+matplotlib = ""
+cartopy = ""
+pandas = ""
+seaborn = ""
+
+[pip.deps]
+pmagpy = "==4.2.106"
+PyQt5 = "" 
+
+[extras]
+channels = ["conda-forge"]

Project.toml

@@ -1,30 +1,30 @@
 name = "SphereUDE"
 uuid = "d7416ba7-148a-4110-b27d-9087fcebab2d"
 authors = ["Facundo Sapienza <fsapienza@berkeley.edu>", "Jordi Bolibar <jordi.bolibar@gmail.com>"]
-version = "0.1.3"
+version = "0.1.4"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DiffEqFlux = "aae7a2af-3d4f-5e19-a356-7da93b79d9d0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
-OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
-PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
-PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -39,22 +39,20 @@ Distributions = "0.25"
 FastGaussQuadrature = "1"
 ForwardDiff = "0.10"
 Infiltrator = "1.2"
-julia = "1.10"
 LineSearches = "7"
 Lux = "1.0"
 Optimization = "3.12, 4"
-OptimizationOptimisers = "0.1.2, 0.3"
 OptimizationOptimJL = "0.1.5, 0.4"
+OptimizationOptimisers = "0.1.2, 0.3"
 OrdinaryDiffEqCore = "1.6.0"
 OrdinaryDiffEqTsit5 = "1.1.0"
 Plots = "1"
 PrettyTables = "2"
-PyCall = "1.9"
-PyPlot = "2.11"
 Revise = "3.1"
 SciMLSensitivity = "7.20"
 Statistics = "1"
 Zygote = "0.6"
+julia = "1.10"
 
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

examples/Torsvik_2012/Gondwana.jl

@@ -51,23 +51,22 @@ data = SphereData(times=times, directions=X, kappas=kappas, L=nothing)
 
 tspan = [times[begin], times[end]]
 
-# params = SphereParameters(tmin = tspan[1], tmax = tspan[2], 
-#                           reg = [Regularization(order=1, power=2.0, λ=1e5, diff_mode=FiniteDifferences(1e-4))], 
-#                           # reg = nothing, 
-#                           pretrain = false, 
-#                           u0 = [0.0, 0.0, -1.0], ωmax = ω₀, 
-#                           reltol = 1e-6, abstol = 1e-6,
-#                           niter_ADAM = 5000, niter_LBFGS = 5000, 
-#                           sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true))) 
 params = SphereParameters(tmin = tspan[1], tmax = tspan[2], 
-                          reg = [Regularization(order=1, power=2.0, λ=1.0, diff_mode=FiniteDifferences(1e-4))], 
+                          reg = [Regularization(order=1, power=2.0, λ=1e5, diff_mode=FiniteDifferences(1e-4))], 
                           # reg = nothing, 
                           pretrain = false, 
                           u0 = [0.0, 0.0, -1.0], ωmax = ω₀, 
                           reltol = 1e-6, abstol = 1e-6,
-                          niter_ADAM = 2000, niter_LBFGS = 2000, 
-                          sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true)), 
-                          hyperparameter_balance = true) 
+                          niter_ADAM = 5000, niter_LBFGS = 5000, 
+                          sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true))) 
+# params = SphereParameters(tmin = tspan[1], tmax = tspan[2], 
+#                           reg = [Regularization(order=1, power=2.0, λ=1.0, diff_mode=FiniteDifferences(1e-4))], 
+#                           pretrain = false, 
+#                           u0 = [0.0, 0.0, -1.0], ωmax = ω₀, 
+#                           reltol = 1e-6, abstol = 1e-6,
+#                           niter_ADAM = 2000, niter_LBFGS = 2000, 
+#                           sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true)), 
+#                           hyperparameter_balance = false) 
 
 
 init_bias(rng, in_dims) = LinRange(tspan[1], tspan[2], in_dims)
@@ -84,7 +83,8 @@ results = train(data, params, rng, nothing, U)
 results_dict = convert2dict(data, results)
 
 
-# JLD2.@save "examples/Torsvik_2012/results/results_dict.jld2" results_dict
+JLD2.@save "examples/Torsvik_2012/results/results_dict.jld2" results_dict
 
 plot_sphere(data, results, -30., 0., saveas="examples/Torsvik_2012/plots/plot_sphere.pdf", title="Double rotation")
+plot_sphere(data, results, -30., 0., saveas="examples/Torsvik_2012/plots/plot_sphere.png", title="Double rotation")
 plot_L(data, results, saveas="examples/Torsvik_2012/plots/plot_L.pdf", title="Double rotation")

examples/Torsvik_2012/Laurentia.jl

@@ -73,6 +73,86 @@ results = train(data, params, rng, nothing, U)
 results_dict = convert2dict(data, results)
 
 
+# JLD2.@save "examples/Torsvik_2012/results/results_dict.jld2" results_dict
+
+plot_sphere(data, results, -30., 0., saveas="examples/Torsvik_2012/plots/Laurentia_plot_sphere.pdf", title="Double rotation")
+plot_L(data, results, saveas="examples/Torsvik_2012/plots/Laurentia_plot_L.pdf", title="Double rotation")
+
+using Pkg; Pkg.activate(".")
+using Revise 
+using Lux
+
+using LinearAlgebra, Statistics, Distributions 
+using SciMLSensitivity
+# using OrdinaryDiffEqCore, OrdinaryDiffEqTsit5
+using Optimization, OptimizationOptimisers, OptimizationOptimJL
+
+using SphereUDE
+
+# Random seed
+using Random
+rng = Random.default_rng()
+Random.seed!(rng, 613)
+
+using DataFrames, CSV
+using Serialization, JLD2
+
+df = CSV.read("./examples/Torsvik_2012/Torsvik-etal-2012_dataset.csv", DataFrame, delim=",")
+
+# Filter the plates that were once part of the supercontinent Gondwana
+
+Laurentia = ["north_america", "greenland"]
+
+df = filter(row -> row.Plate ∈ Laurentia, df)
+df.Times = df.Age .+= rand(sampler(Normal(0,0.1)), nrow(df))  # Needs to fix this! 
+
+df = sort(df, :Times)
+times = df.Times
+
+# Fill missing values
+df.RLat .= coalesce.(df.RLat, df.Lat)
+df.RLon .= coalesce.(df.RLon, df.Lon)
+
+X = sph2cart(Matrix(df[:,["RLat","RLon"]])'; radians=false)
+
+# Retrieve uncertanties from poles and convert α95 into κ
+kappas = (140.0 ./ df.a95).^2
+
+data = SphereData(times=times, directions=X, kappas=kappas, L=nothing)
+
+# Training
+
+# Expected maximum angular deviation in one unit of time (degrees)
+Δω₀ = 1.5   
+# Angular velocity 
+ω₀ = Δω₀ * π / 180.0
+
+tspan = [times[begin], times[end]]
+
+params = SphereParameters(tmin = tspan[1], tmax = tspan[2], 
+                          reg = [Regularization(order=1, power=2.0, λ=1e6, diff_mode=FiniteDifferences(1e-4))], 
+                        #   reg = nothing, 
+                          pretrain = false, 
+                          u0 = [0.0, 0.0, -1.0], ωmax = ω₀, 
+                          reltol = 1e-6, abstol = 1e-6,
+                          niter_ADAM = 5000, niter_LBFGS = 5000, 
+                          sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true))) 
+
+
+init_bias(rng, in_dims) = LinRange(tspan[1], tspan[2], in_dims)
+init_weight(rng, out_dims, in_dims) = 0.1 * ones(out_dims, in_dims)
+
+# Customized neural network to similate weighted moving window in L
+U = Lux.Chain(
+    Lux.Dense(1, 200, rbf, init_bias=init_bias, init_weight=init_weight, use_bias=true),
+    Lux.Dense(200,10, gelu),
+    Lux.Dense(10, 3, Base.Fix2(sigmoid_cap, params.ωmax), use_bias=false)
+)
+
+results = train(data, params, rng, nothing, U)
+results_dict = convert2dict(data, results)
+
+
 # JLD2.@save "examples/Torsvik_2012/results/results_dict.jld2" results_dict
 
 plot_sphere(data, results, -30., 0., saveas="examples/Torsvik_2012/plots/Laurentia_plot_sphere.pdf", title="Double rotation")

examples/Torsvik_2012/plots/plot_results.jl

@@ -7,6 +7,23 @@ using Plots.PlotMeasures
 using JLD2
 using LinearAlgebra
 
+blue_belize = RGBA(41/255, 128/255, 185/255, 1)
+blue_midnigth = RGBA(44/255, 62/255, 80/255, 1)
+purple_wisteria = RGBA(142/255, 68/255, 173/255, 1)
+red_pomegrade = RGBA(192/255, 57/255, 43/255, 1)
+orange_carrot = RGBA(230/255, 126/255, 34/255, 1)
+green_nephritis = RGBA(39/255, 174/255, 96/255, 1)
+green_sea = RGBA(22/255, 160/255, 133/255, 1)
+
+# colors 
+lat_color_scatter = blue_belize
+lat_color_line = purple_wisteria
+lon_color_scatter = red_pomegrade
+lon_color_line = orange_carrot
+angular_color_scatter = green_nephritis
+angular_color_line = green_sea
+loss_color = blue_midnigth
+
 JLD2.@load "examples/Torsvik_2012/results/results_dict.jld2" 
 
 data_directions_sph = cart2sph(results_dict["directions"], radians=false)
@@ -22,10 +39,10 @@ fit_lon  = fit_directions_sph[2,:]
 
 # Latitude Plot 
 
-Plots.scatter(results_dict["times"], data_lat, label="Paleopole Latitudes", yerr=α_error, c=:lightsteelblue2, markersize=5)
+Plots.scatter(results_dict["times"], data_lat, label="Paleopole Latitudes", yerr=α_error, c=lat_color_scatter, ms=5, msw=0.3)
 plot_latitude = Plots.plot!(results_dict["fit_times"], fit_lat, label="Estimated APWP using SphereUDE", 
-                    xlabel="Age (Myr)", yticks=[-90, -60, -30, 0, 30, 60],
-                    ylabel="Latitude (degrees)", ylims=(-90,60), lw = 4, c=:brown,
+                    xlabel="Age (Myr)", yticks=[-90, -60, -30, 0, 30, 60], xlims=(0,520),
+                    ylabel="Latitude (degrees)", ylims=(-90,60), lw = 4, c=lat_color_line, 
                     legend=:topleft)
 plot!(fontfamily="Computer Modern",
      #title="PIL51",
@@ -46,10 +63,10 @@ Plots.savefig(plot_latitude, "examples/Torsvik_2012/plots/latitude.pdf")
 
 α_error_lon = α_error ./ cos.(π ./ 180. .* data_lat)
 
-Plots.scatter(results_dict["times"], data_lon, label="Paleopole Longitudes", yerr=α_error_lon, c=:lightsteelblue2, markersize=5)
+Plots.scatter(results_dict["times"], data_lon, label="Paleopole Longitudes", yerr=α_error_lon, c=lon_color_scatter, markersize=5, msw=0.3)
 plot_longitude = Plots.plot!(results_dict["fit_times"], fit_lon, label="Estimated APWP using SphereUDE", 
-                    xlabel="Age (Myr)", yticks=[-180, -120, -60 , 0, 60, 120, 180],
-                    ylabel="Longitude (degrees)", ylims=(-180,180), lw = 4, c=:brown,
+                    xlabel="Age (Myr)", yticks=[-180, -90 , 0, 90, 180], xlims=(0,520),
+                    ylabel="Longitude (degrees)", ylims=(-220,180), lw = 4, c=lon_color_line,
                     legend=:bottomright)
 plot!(fontfamily="Computer Modern",
      #title="PIL51",
@@ -65,33 +82,17 @@ plot!(fontfamily="Computer Modern",
 
 Plots.savefig(plot_longitude, "examples/Torsvik_2012/plots/longitude.pdf")
 
-### Lat and long combined
-
-combo_plot = plot(plot_latitude, plot_longitude, layout = (2, 1))
-plot!(fontfamily="Computer Modern",
-     #title="PIL51",
-    titlefontsize=18,
-    tickfontsize=15,
-    legendfontsize=15,
-    guidefontsize=18,
-    #ylimits=(0.1,10),
-    #xlimits=(10^(-4),10^(-1)),
-    margin= 7mm,
-    size=(1200,800),
-    dpi=600)
-Plots.savefig(combo_plot, "examples/Torsvik_2012/plots/latitude_longitude.pdf")
-
 
 ### Angular velocity Plot
 
 angular_velocity = mapslices(x -> norm(x), results_dict["fit_rotations"], dims=1)[:]
 angular_velocity_path = [norm(cross(results_dict["fit_directions"][:,i], results_dict["fit_rotations"][:,i] )) for i in axes(results_dict["fit_rotations"],2)]
 
 plot_angular = Plots.plot(results_dict["fit_times"], angular_velocity, label="Maximum total angular velocity", 
-                    xlabel="Age (Myr)", 
-                    ylabel="Angular velocity (degrees/My)", lw = 5, c=:darkgreen,
+                    xlabel="Age (Myr)", yticks=[0.0, 0.01, 0.02, 0.03, 0.04], xlims=(0,520), ylims=(0.0, 0.041),
+                    ylabel="Angular velocity (degrees/My)", lw = 5, c=angular_color_scatter,
                     legend=:topleft)
-plot!(results_dict["fit_times"], angular_velocity_path, label="Pole angular velocity", lw = 4, c=:lightseagreen, ls=:dot)
+plot!(results_dict["fit_times"], angular_velocity_path, label="Pole angular velocity", lw = 4, c=angular_color_line, ls=:dot)
 plot!(fontfamily="Computer Modern",
     #title="PIL51",
     titlefontsize=18,
@@ -100,21 +101,35 @@ plot!(fontfamily="Computer Modern",
     guidefontsize=18,
     #ylimits=(0.1,10),
     #xlimits=(10^(-4),10^(-1)),
-    margin= 10mm,
+    margin= 7mm,
     size=(1200,500),
-    dpi=300)
+    dpi=600)
 
 
 Plots.savefig(plot_angular, "examples/Torsvik_2012/plots/angular.pdf")
 
+### Lat and long combined
+
+combo_plot = plot(plot_latitude, plot_longitude, plot_angular, layout = (3, 1))
+plot!(fontfamily="Computer Modern",
+     #title="PIL51",
+    # titlefontsize=18,
+    # tickfontsize=15,
+    # legendfontsize=15,
+    # guidefontsize=18,
+    #ylimits=(0.1,10),
+    #xlimits=(10^(-4),10^(-1)),
+    # margin= 7mm,
+    size=(1200,1000))
+Plots.savefig(combo_plot, "examples/Torsvik_2012/plots/latitude_longitude.pdf")
 
 ### Loss function
 
 losses = results_dict["losses"]
 
 plot_loss = Plots.plot(1:length(losses), losses, label="Loss Function", 
                     xlabel="Epoch", 
-                    ylabel="Loss", lw = 5, c=:indigo,
+                    ylabel="Loss", lw = 5, c=loss_color,
                     yaxis=:log,
                     # yticks=[1,10,100],
                     xlimits=(0,10000),