From 0c189c577adcba2847018174f3d43bb382fa1e0c Mon Sep 17 00:00:00 2001
From: islent <leoislent@gmail.com>
Date: Tue, 13 Aug 2024 15:54:27 +0800
Subject: [PATCH] fix Rotation Curve

---
 Project.toml   |  2 ++
 src/AstroIC.jl |  1 +
 src/disk.jl    | 34 +++++++++++++++++++++++++++++-----
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8783f17..87dde27 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,6 +9,7 @@ AstroLib = "c7932e45-9af1-51e7-9da9-f004cd3a462b"
 AstroSimBase = "c6a34d98-f626-4d8d-b450-a3f124eaada6"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
@@ -26,6 +27,7 @@ AstroIO = "0.1"
 AstroLib = "0.4"
 AstroSimBase = "0.1"
 BangBang = "0.3, 0.4"
+Dierckx = "0.5"
 Distributions = "0.25"
 DocStringExtensions = "0.8, 0.9"
 PhysicalConstants = "0.2"
diff --git a/src/AstroIC.jl b/src/AstroIC.jl
index a44088d..d9939df 100644
--- a/src/AstroIC.jl
+++ b/src/AstroIC.jl
@@ -11,6 +11,7 @@ using Random
 using BangBang
 using Dates
 using StructArrays
+using Dierckx
 
 @reexport using PhysicalParticles
 @reexport using AstroIO
diff --git a/src/disk.jl b/src/disk.jl
index 2820756..13a29ea 100644
--- a/src/disk.jl
+++ b/src/disk.jl
@@ -116,6 +116,30 @@ function sech2_cdf_inv(u)
     return atanh(2*u-1)
 end
 
+"""
+$(TYPEDSIGNATURES)
+Rotate along the positive z-axis.
+- `ratio`: rotational motion v.s. random motion. If equals `1`, only rotational component
+"""
+function rotational_velocity_acc(x, y, z, a, ratio = 1.0)
+    u = unit(x)
+    r = sqrt(x^2 + y^2 + z^2)
+    v = sqrt(a * r)
+    v_vec = -normalize(PVector(ustrip(u,x), ustrip(u,y), 0.0) × PVector(0.0, 0.0, 1.0)) * v
+    v_vec = ratio * v_vec + (1-ratio) * randn(PVector{Float64}) * v
+end
+
+"""
+$(TYPEDSIGNATURES)
+Rotate along the positive z-axis.
+- `ratio`: rotational motion v.s. random motion. If equals `1`, only rotational component
+"""
+function rotational_velocity(x, y, v, ratio = 1.0)
+    u = unit(x)
+    v_vec = -normalize(PVector(ustrip(u,x), ustrip(u,y), 0.0) × PVector(0.0, 0.0, 1.0)) * v
+    v_vec = ratio * v_vec + (1-ratio) * randn(PVector{Float64}) * v
+end
+
 """
 $(TYPEDSIGNATURES)
 
@@ -124,9 +148,9 @@ function generate(config::ExponentialDisk, units = uAstro;
         RotationCurve = nothing,
         MaxRadius = 5 * config.ScaleRadius,
         MaxHeight = MaxRadius,
-        k = 1, # degree of rotation curve interpolation/extrapolation spline (1 = linear, 2 = quadratic, 3 = cubic, up to 5)
+        k = 2, # degree of rotation curve interpolation/extrapolation spline (1 = linear, 2 = quadratic, 3 = cubic, up to 5)
         bc = "nearest", # behavior when evaluating the spline outside the support domain, which is (minimum(x), maximum(x)). The allowed values are "nearest", "zero", "extrapolate", "error"
-        rotational_ratio = 0,
+        rotational_ratio = 0.9,
     )
     uLen = getuLength(units)
     uVel = getuVel(units)
@@ -184,9 +208,9 @@ function generate(config::ExponentialDisk, units = uAstro;
         vel = [PVector(uVel) for i in 1:NumSamples]
     else
         xc, vc = RotationCurve
-        spl = Spline1D(ustrip.(unit(eltype(xc)), xc), vc; k, bc)
-        v = spl(ustrip.(unit(eltype(R)), R))
-        vel = rotational_velocity.(pos.x, pos.y, pos.z, v, rotational_ratio)
+        spl = Spline1D(ustrip.(uLen, xc), ustrip.(uVel, vc); k, bc)
+        v = spl(ustrip.(uLen, R)) * uVel
+        vel = rotational_velocity.(pos.x, pos.y, v, rotational_ratio)
     end
 
     # Packing