Skip to content

Commit

Permalink
Merge pull request #172 from RemoteSensingTools/GPU-library-update
Browse files Browse the repository at this point in the history
Fixed GPU breaking changes in new tools
  • Loading branch information
cfranken authored Jun 17, 2024
2 parents bee803a + 486fa31 commit e385c0e
Show file tree
Hide file tree
Showing 13 changed files with 63 additions and 55 deletions.
4 changes: 2 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "vSmartMOM"
uuid = "7ba11eeb-0a61-4a04-a413-bf612cc2007e"
authors = ["Christian Frankenberg <[email protected]>, Suniti Sanghavi ([email protected]), Rupesj Jeyaram and contributors"]
version = "1.0.2"
authors = ["Suniti Sanghavi ([email protected]), Christian Frankenberg <[email protected]>, Rupesh Jeyaram and contributors"]
version = "1.0.3"

[deps]
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
Expand Down
4 changes: 2 additions & 2 deletions src/Architectures.jl
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ macro hascuda(expr)
end

devi(::CPU) = KernelAbstractions.CPU()
devi(::GPU) = CUDAKernels.CUDADevice()
devi(::GPU) = CUDA.CUDABackend(; always_inline=true)

architecture(::Array) = CPU()
@hascuda architecture(::CuArray) = GPU()
Expand All @@ -52,6 +52,6 @@ devi(::GPU) = CUDAKernels.CUDADevice()

default_architecture = has_cuda() ? GPU() : CPU()

synchronize_if_gpu() = has_cuda() ? synchronize() : nothing
synchronize_if_gpu() = has_cuda() ? CUDA.synchronize() : nothing

end
3 changes: 3 additions & 0 deletions src/CoreRT/CoreKernel/interaction.jl
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,9 @@ function interaction_helper!(::ScatteringInterface_11, SFI,
# Repeating for mirror-reflected directions

# Compute and store `(I - r⁻⁺ * R⁺⁻)⁻¹`
#handle = CUBLAS.handle()
#CUBLAS.math_mode!(handle, CUDA.FAST_MATH)
#@show typeof(I_static .- R⁺⁻ ⊠ r⁻⁺)
@timeit "interaction inv2" batch_inv!(tmp_inv, I_static .- R⁺⁻ r⁻⁺)
# T₂₁(I-R₀₁R₂₁)⁻¹
T21_inv = t⁺⁺ tmp_inv
Expand Down
14 changes: 7 additions & 7 deletions src/CoreRT/CoreKernel/interaction_hdrf.jl
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,10 @@ function interaction_hdrf!(SFI,
NquadN = Nquad * pol_type.n
hdr_J₀⁻ .= r⁻⁺ J₀⁺ .+ j₀⁻
# @show hdr_J₀⁻./ J₀⁺
@show hdr_J₀⁻[1,1,:]
@show J₀⁺[1,1,:]
@show iμ₀
@show j₀⁺[iμ₀Nstart,1,:]
#@show hdr_J₀⁻[1,1,:]
#@show J₀⁺[1,1,:]
#@show iμ₀
#@show j₀⁺[iμ₀Nstart,1,:]
qp = Array(qp_μN)
if m==0

Expand All @@ -32,9 +32,9 @@ function interaction_hdrf!(SFI,
j=i:pol_type.n:NquadN
bhr_J₀⁻[i,:] .= Array(sum(hdr_J₀⁻[j,1,:].*wt_μN[j].*qp_μN[j], dims=1)')
bhr_J₀⁺[i,:] .= Array(sum(J₀⁺[j,1,:].*wt_μN[j].*qp_μN[j], dims=1)' .+ j₀⁺[iμ₀Nstart,1,:] .* qp[iμ₀Nstart])
if i==1
@show bhr_J₀⁻[i,:], bhr_J₀⁺[i,:]
end
#if i==1
# @show bhr_J₀⁻[i,:], bhr_J₀⁺[i,:]
#end
#TODO: Use Radau quadrature and include insolation in the quadrature sum

#@show j₀⁺[iμ₀,1,1:3].* qp[iμ₀], J₀⁺[iμ₀,1,1:3] .* qp[iμ₀], bhr_J₀⁺[i,1:3], bhr_J₀⁻[i,1:3]
Expand Down
9 changes: 5 additions & 4 deletions src/CoreRT/atmo_prof.jl
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ function read_atmos_profile(file_path::String)
vmr = convert(Dict{String, Union{Real, Vector}}, params_dict["vmr"])

# Return the atmospheric profile struct
return AtmosphericProfile(T, q, p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o, vmr)
return AtmosphericProfile(T, q, p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o, vmr,Δz)

end

Expand All @@ -105,14 +105,15 @@ function reduce_profile(n::Int, profile::AtmosphericProfile{FT}) where {FT}
@assert n < length(profile.T)

# Unpack the profile vmr
@unpack vmr = profile
@unpack vmr, Δz = profile

# New rough half levels (boundary points)
a = range(0, maximum(profile.p_half), length=n+1)

# Matrices to hold new values
T = zeros(FT, n);
q = zeros(FT, n);
Δz_ = zeros(FT, n);
p_full = zeros(FT, n);
p_half = zeros(FT, n+1);
vmr_h2o = zeros(FT, n);
Expand All @@ -138,7 +139,7 @@ function reduce_profile(n::Int, profile::AtmosphericProfile{FT}) where {FT}
p_full[i] = mean(profile.p_full[ind])
T[i] = mean(profile.T[ind])
q[i] = mean(profile.q[ind])

Δz_[i] = sum(Δz[ind])
vcd_dry[i] = sum(profile.vcd_dry[ind])
vcd_h2o[i] = sum(profile.vcd_h2o[ind])
vmr_h2o[i] = vcd_h2o[i]/vcd_dry[i]
Expand All @@ -157,7 +158,7 @@ function reduce_profile(n::Int, profile::AtmosphericProfile{FT}) where {FT}
end
end

return AtmosphericProfile(T, p_full, q, p_half, vmr_h2o, vcd_dry, vcd_h2o, new_vmr)
return AtmosphericProfile(T, p_full, q, p_half, vmr_h2o, vcd_dry, vcd_h2o, new_vmr,Δz_)
end

"""
Expand Down
3 changes: 3 additions & 0 deletions src/CoreRT/gpu_batched.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ This file contains implementations of batched linear algebra code
=#

@inline synchronize() = CUDA.synchronize()

"Given 3D CuArrays A and B, fill in X[:,:,k] = A[:,:,k] \\ B[:,:,k]"
function batch_solve!(X::CuArray{FT,3}, A::CuArray{FT,3}, B::CuArray{FT,3}) where {FT}

Expand Down Expand Up @@ -32,6 +34,7 @@ end

"Given 3D CuArray A, fill in X[:,:,k] = A[:,:,k] \\ I"
function batch_inv!(X::CuArray{FT,3}, A::CuArray{FT,3}) where {FT}
#CUBLAS.math_mode!(CUBLAS.handle(), CUDA.FAST_MATH)
# LU-factorize A
pivot, info = CUBLAS.getrf_strided_batched!(A, true);synchronize()
# Invert LU factorization of A
Expand Down
19 changes: 9 additions & 10 deletions src/CoreRT/model_from_parameters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ function model_from_parameters(params::vSmartMOM_Parameters)
# This is a kludge for now, tau_abs sometimes needs to be a dual. Suniti & us need to rethink this all!!
# i.e. code the rt core with fixed amount of derivatives as in her paper, then compute chain rule for dtau/dVMr, etc...
FT2 = isnothing(params.absorption_params) || !haskey(params.absorption_params.vmr,"CO2") ? params.float_type : eltype(params.absorption_params.vmr["CO2"])
τ_abs = [zeros(FT2, length(params.spec_bands[i]), length(profile.p_full)) for i in 1:n_bands]
τ_abs = [zeros(FT, length(params.spec_bands[i]), length(profile.p_full)) for i in 1:n_bands]

# Loop over all bands:
for i_band=1:n_bands
Expand All @@ -54,7 +54,7 @@ function model_from_parameters(params::vSmartMOM_Parameters)
curr_band_λ = FT.(1e4 ./ params.spec_bands[i_band])
# @show profile.vcd_dry, size(τ_rayl[i_band])
# Compute Rayleigh properties per layer for `i_band` band center

#@show τ_rayl[i_band]
τ_rayl[i_band] .= getRayleighLayerOptProp(profile.p_half[end],
curr_band_λ, #(mean(curr_band_λ)),
params.depol, profile.vcd_dry);
Expand Down Expand Up @@ -96,7 +96,8 @@ function model_from_parameters(params::vSmartMOM_Parameters)
#FT2 = params.float_type

# τ_aer[iBand][iAer,iZ]
τ_aer = [zeros(FT2, n_aer, length(profile.p_full)) for i=1:n_bands];
# Again, be careful with Dual Numbers
τ_aer = [zeros(FT, n_aer, length(profile.p_full)) for i=1:n_bands];

# Loop over aerosol type
for i_aer=1:n_aer
Expand All @@ -123,10 +124,9 @@ function model_from_parameters(params::vSmartMOM_Parameters)
params.scattering_params.nquad_radius)
mie_model.aerosol.nᵣ = real(params.scattering_params.n_ref)
mie_model.aerosol.nᵢ = -imag(params.scattering_params.n_ref)
@show params.scattering_params.n_ref
# k for reference wavelength
k_ref = compute_ref_aerosol_extinction(mie_model, params.float_type)
@show k_ref
#params.scattering_params.rt_aerosols[i_aer].p₀, params.scattering_params.rt_aerosols[i_aer].σp

# Loop over bands
for i_band=1:n_bands

Expand All @@ -143,9 +143,9 @@ function model_from_parameters(params::vSmartMOM_Parameters)
params.scattering_params.nquad_radius)
n_ref = params.scattering_params.n_ref
k = compute_ref_aerosol_extinction(mie_model, params.float_type)
@show k

#@show k
# Compute raw (not truncated) aerosol optical properties (not needed in RT eventually)
#@show FT2, FT
@timeit "Mie calc" aerosol_optics_raw = compute_aerosol_optical_properties(mie_model, FT);
@show aerosol_optics_raw.k
# Compute truncated aerosol optical properties (phase function and fᵗ), consistent with Ltrunc:
Expand Down Expand Up @@ -246,9 +246,8 @@ function model_from_parameters(RS_type::Union{VS_0to1_plus, VS_1to0_plus},
greek_rayleigh = Scattering.get_greek_rayleigh(params.depol)
τ_rayl = [zeros(params.float_type,1, length(profile.p_full)) for i=1:n_bands];

# τ_abs[iBand][iSpec,iZ]
# Pre-allocated absorption arrays
τ_abs = [zeros(params.float_type, length(params.spec_bands[i]), length(profile.p_full)) for i in 1:n_bands]
@show params.absorption_params.molecules[2]
# Loop over all bands:
for i_band=1:n_bands
@show params.spec_bands[i_band]
Expand Down
2 changes: 1 addition & 1 deletion src/CoreRT/parameters_from_yaml.jl
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ function parameters_from_yaml(file_path)
if !haskey(params_dict["scattering"],"n_ref")
n_ref = aerosols[1].aerosol.nᵣ - im*aerosols[1].aerosol.nᵢ
else
@show params_dict["scattering"]["n_ref"]
#@show params_dict["scattering"]["n_ref"]
n_ref = params_dict["scattering"]["n_ref"]
end
scattering_params = ScatteringParameters(aerosols, r_max, nquad_radius,
Expand Down
4 changes: 2 additions & 2 deletions src/CoreRT/postprocessing_vza_ms.jl
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ function postprocessing_vza_ms!(RS_type::Union{RRS, VS_0to1_plus, VS_1to0_plus},

tuwJ = [zeros(typeof(topJ₀⁺[1][1,1,1]), (length(topJ₀⁺[1][:,1,1]), 1, nSpec)) for i=1:length(sensor_levels)] #similar(topJ₀⁺); #deepcopy(topJ₀⁺)
tdwJ = [zeros(typeof(topJ₀⁺[1][1,1,1]), (length(topJ₀⁺[1][:,1,1]), 1, nSpec)) for i=1:length(sensor_levels)]#similar(topJ₀⁺); #deepcopy(topJ₀⁺)
@show size(tuwJ), size(tdwJ)
#@show size(tuwJ), size(tdwJ)
botieR⁻⁺ = (composite_layer.botieR⁻⁺);
topieR⁺⁻ = (composite_layer.topieR⁺⁻);

Expand All @@ -116,7 +116,7 @@ function postprocessing_vza_ms!(RS_type::Union{RRS, VS_0to1_plus, VS_1to0_plus},
#tdwieJ = deepcopy(topieJ₀⁺)
tuwieJ = [zeros(typeof(topJ₀⁺[1][1,1,1]), (length(topJ₀⁺[1][:,1,1]), 1, nSpec, length(topieJ₀⁺[1][1,1,1,:]))) for i=1:length(sensor_levels)] #similar(topJ₀⁺); #deepcopy(topJ₀⁺)
tdwieJ = [zeros(typeof(topJ₀⁺[1][1,1,1]), (length(topJ₀⁺[1][:,1,1]), 1, nSpec, length(topieJ₀⁺[1][1,1,1,:]))) for i=1:length(sensor_levels)] #similar(topJ₀⁺); #deepcopy(topJ₀⁺)
@show size(tuwieJ), size(tdwieJ)
#@show size(tuwieJ), size(tdwieJ)
for ims = 1:length(sensor_levels)
if(sensor_levels[ims]==0)
# elastic
Expand Down
13 changes: 7 additions & 6 deletions src/Scattering/compute_NAI2.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ function compute_aerosol_optical_properties(model::MieModel{FDT}, FT2::Type=Floa
#@show size_distribution.σ
# TODO: This is still very clumsy, the FT conversions are not good here.
FT = eltype(nᵣ);

@show "Mie", FT, FT2
#@show FT, ForwardDiff.valtype(size_distribution.σ)
vFT = ForwardDiff.valtype(nᵣ)
#@assert FT == Float64 "Aerosol computations require 64bit"
Expand Down Expand Up @@ -161,7 +161,7 @@ function compute_aerosol_optical_properties(model::MieModel{FDT}, FT2::Type=Floa

# Check whether this is a Dual number (if so, don't do any conversions)
# TODO: Equally clumsy, needs to be fixed.
if FT <: AbstractFloat
#if FT <: AbstractFloat
#@show "Convert greek", FT2
# Create GreekCoefs object with α, β, γ, δ, ϵ, ζ
greek_coefs = GreekCoefs(convert.(FT2, α),
Expand All @@ -172,12 +172,13 @@ function compute_aerosol_optical_properties(model::MieModel{FDT}, FT2::Type=Floa
convert.(FT2, ζ))
#@show typeof(convert.(FT2, β)), typeof(greek_coefs)
# Return the packaged AerosolOptics object
@show greek_coefs
return AerosolOptics(greek_coefs=greek_coefs, ω̃=FT2(bulk_C_sca / bulk_C_ext), k=FT2(bulk_C_ext), fᵗ=FT2(1))

else
greek_coefs = GreekCoefs(α, β, γ,δ,ϵ,ζ)
return AerosolOptics(greek_coefs=greek_coefs, ω̃=(bulk_C_sca / bulk_C_ext), k=(bulk_C_ext), fᵗ=FT(1))
end
#else
# greek_coefs = GreekCoefs(α, β, γ,δ,ϵ,ζ)
# return AerosolOptics(greek_coefs=greek_coefs, ω̃=(bulk_C_sca / bulk_C_ext), k=(bulk_C_ext), fᵗ=FT(1))
#end
end

function compute_ref_aerosol_extinction(model::MieModel{FDT}, FT2::Type=Float64) where FDT <: NAI2
Expand Down
15 changes: 8 additions & 7 deletions src/Scattering/truncate_phase.jl
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,8 @@ function truncate_phase(mod::δBGE, aero::AerosolOptics{FT}; reportFit=false) wh
l_tr = l_max
# Obtain Gauss-Legendre quadrature points and weights for phase function:
μ, w_μ = gausslegendre(length(β));

μ = convert.(FT, μ); w_μ = convert.(FT, w_μ);
#show μ
# Reconstruct phase matrix elements:
scattering_matrix, P, P² = reconstruct_phase(greek_coefs, μ; returnLeg=true)

Expand All @@ -125,9 +126,9 @@ function truncate_phase(mod::δBGE, aero::AerosolOptics{FT}; reportFit=false) wh
Aᵢⱼ = ∑ₖ w_μₖ Pᵢ(μₖ)Pⱼ(μₖ)/f₁₁²(μₖ), xᵢ=cᵢ (as in Sanghavi & Stephens 2015), and
bᵢ = ∑ₖ w_μₖ Pᵢ(μₖ)/f₁₁(μₖ)
=#
A = zeros(l_tr, l_tr)
x = zeros(l_tr)
b = zeros(l_tr)
A = zeros(FT,l_tr, l_tr)
x = zeros(FT,l_tr)
b = zeros(FT,l_tr)

for i = 1:l_tr
b[i] = sum(w_μ.*P[:,i]./f₁₁)
Expand All @@ -150,9 +151,9 @@ function truncate_phase(mod::δBGE, aero::AerosolOptics{FT}; reportFit=false) wh
Aᵢⱼ = ∑ₖ w_μₖ facᵢP²ᵢ(μₖ)facⱼP²ⱼ(μₖ)/f₁₂²(μₖ), xᵢ=gᵢ (as in Sanghavi & Stephens 2015), and
bᵢ = ∑ₖ w_μₖ facᵢP²ᵢ(μₖ)/f₁₂(μₖ)
=#
A = zeros(l_tr, l_tr)
x = zeros(l_tr)
b = zeros(l_tr)
A = zeros(FT,l_tr, l_tr)
x = zeros(FT,l_tr)
b = zeros(FT,l_tr)

for i = 3:l_tr
b[i] = fac[i]*sum(w_μ.*P²[:,i]./f₁₂)
Expand Down
16 changes: 8 additions & 8 deletions test/prototyping/AD_OCO2_africa.jl
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,14 @@ using LinearAlgebra
# Load parameters from file
parameters = parameters_from_yaml("test/test_parameters/3BandParameters.yaml")
#parameters.architecture = CPU()
FT = Float32
FT = Float64
parameters.float_type = FT
# Load OCO Data:
# File names:

# Africa
L1File = "/net/fluo/data1/group/oco2/L1bSc/oco2_L1bScND_31182a_200512_B10206r_210427151019.h5"
metFile = "/net/fluo/data1/group/oco2/L2Met/oco2_L2MetND_31182a_200512_B10206r_210425233338.h5"
L1File = "/net/fluo/data3/data/FluoData1/group/oco2/L1bSc/oco2_L1bScND_31182a_200512_B10206r_210427151019.h5"
metFile = "/net/fluo/data3/data/FluoData1/group/oco2/L2Met/oco2_L2MetND_31182a_200512_B10206r_210425233338.h5"
dictFile = "/home/cfranken/code/gitHub/InstrumentOperator.jl/json/oco2.yaml"
# Load L1 file (could just use filenames here as well)
oco = InstrumentOperator.load_L1(dictFile,L1File, metFile);
Expand Down Expand Up @@ -82,8 +82,8 @@ function runner!(y, x, parameters=parameters, oco_sounding= oco_sounding, Tsolar
CoreRT.LambertianSurfaceLegendre([x[4],x[5],x[6]]),
CoreRT.LambertianSurfaceLegendre([x[7],x[8],x[9]])];

parameters.scattering_params.rt_aerosols[1].τ_ref = exp(x[10]);
parameters.scattering_params.rt_aerosols[1].p₀ = x[11]
parameters.scattering_params.rt_aerosols[1].τ_ref = x[10];
#parameters.scattering_params.rt_aerosols[1].profile = Normal(x[11],50.0)
size_distribution = LogNormal(x[12], FT(0.3))
parameters.scattering_params.rt_aerosols[1].aerosol.size_distribution = size_distribution
#@show size_distribution
Expand Down Expand Up @@ -149,7 +149,7 @@ function runner2(x, parameters=parameters, oco_sounding= oco_sounding, Tsolar =
CoreRT.LambertianSurfaceLegendre([x[7],x[8],x[9]])];

parameters.scattering_params.rt_aerosols[1].τ_ref = exp(x[10]);
parameters.scattering_params.rt_aerosols[1].p₀ = x[11]
#parameters.scattering_params.rt_aerosols[1].profile.μ = x[11]
size_distribution = LogNormal(x[12], FT(0.3))
parameters.scattering_params.rt_aerosols[1].aerosol.size_distribution = size_distribution
#@show size_distribution
Expand Down Expand Up @@ -217,9 +217,9 @@ xₐ = FT[ 0.376, # Albedo band 1, degree 1
0.1,#0.0512, # Albedo band 3, degree 1
0, # Albedo band 3, degree 2
0.00, # Albedo band 3, degree 3
-1.0, # log AOD
0.02, # log AOD
800.0, # Aerosol peak height (hPa)
-1.0, # log of size distribution mean (1µm here)
1.0, # log of size distribution mean (1µm here)
1.3, # Real refractive index of aerosol
1, # H2O scaling factor for entire profile
400, # CO2 Top layers
Expand Down
12 changes: 6 additions & 6 deletions test/test_parameters/3BandParameters.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ radiative_transfer:
# Depolarization factor
depol: 0.03
# Floating point type for calculations (Float32, Float64)
float_type: Float32
float_type: Float64
# Architecture (default_architecture, GPU(), CPU())
architecture: default_architecture #default_architecture
architecture: CPU() #default_architecture

# =================================================================
# Simulation Geometry
Expand Down Expand Up @@ -63,7 +63,7 @@ atmospheric_profile:
176.85, 236.64, 314.58, 418.87, 557.76, 735.00,
800.12, 849.00, 912.00, 980.00, 1005.0]
# Reduce profile to n layers
profile_reduction: 10
profile_reduction: -1

# =================================================================
# Absorption-Related Parameters (Optional)
Expand All @@ -76,9 +76,9 @@ absorption:
- [H2O,CO2] # Molecules in Band #3
# LookUpTable files (Interpolation Model saved as JLD2!)
LUTfiles:
- ["/net/fluo/data2/data/ABSCO_CS_Database/v5.2_final/o2_v52.jld2"]
- ["/net/fluo/data2/data/ABSCO_CS_Database/v5.2_final/wh2o_v52.jld2", "/net/fluo/data2/data/ABSCO_CS_Database/v5.2_final/wco2_v52.jld2"]
- ["/net/fluo/data2/data/ABSCO_CS_Database/v5.2_final/sh2o_v52.jld2", "/net/fluo/data2/data/ABSCO_CS_Database/v5.2_final/sco2_v52.jld2"]
- ["/net/fluo/data3/data/Databases/ABSCO_CS_Database/v5.2_final/o2_v52.jld2"]
- ["/net/fluo/data3/data/Databases/ABSCO_CS_Database/v5.2_final/wh2o_v52.jld2", "/net/fluo/data3/data/Databases/ABSCO_CS_Database/v5.2_final/wco2_v52.jld2"]
- ["/net/fluo/data3/data/Databases/ABSCO_CS_Database/v5.2_final/sh2o_v52.jld2", "/net/fluo/data3/data/Databases/ABSCO_CS_Database/v5.2_final/sco2_v52.jld2"]
# VMR profiles can either be real-valued numbers,
# or an array of nodal points from TOA to BOA, interpolated in pressure space
vmr:
Expand Down

0 comments on commit e385c0e

Please sign in to comment.