From 1a500a37baf243188ec81f15c000c7aa971c7dcf Mon Sep 17 00:00:00 2001
From: Gaspard Jankowiak <gaspard.jankowiak@uni-graz.at>
Date: Thu, 16 May 2024 14:42:48 +0200
Subject: [PATCH] refactor

---
 PLUV.jl  | 23 ----------------------
 Utils.jl | 59 ++++++++++++++++++++++++++++++++++++++++++++------------
 2 files changed, 47 insertions(+), 35 deletions(-)

diff --git a/PLUV.jl b/PLUV.jl
index 7763872..a1c8a0b 100644
--- a/PLUV.jl
+++ b/PLUV.jl
@@ -325,29 +325,6 @@ function negative_water(P::AbstractArray{Float64})
   return count(<(-1e-3), BW)
 end
 
-function get_free_parameters_idx(metadata::Dict)
-  m = metadata["parameters"]
-  return filter(i -> !m[i]["fixed"], eachindex(m))
-end
-
-function reduce_to_free_parameters(metadata::Dict, params_init::NamedTuple, lb, ub, q)
-  default_params = Float64[params_init...]
-  free_idx = get_free_parameters_idx(metadata)
-  P_reduced = default_params[free_idx]
-  lb_reduced = lb[free_idx]
-  ub_reduced = ub[free_idx]
-  function reduced(P::AbstractArray{Float64})
-    P_reduced .= P
-    return intensity(default_params, q)
-  end
-  return reduced, P_reduced, lb_reduced, ub_reduced
-end
-
-function reduce_to_free_parameters(metadata::Dict, params::Vector{Float64})
-  free_idx = get_free_parameters_idx(metadata)
-  return params[free_idx]
-end
-
 end # module PLUV
 
 # vim ts=4,sts=4,sw=4
diff --git a/Utils.jl b/Utils.jl
index 379b1c8..4f31ade 100644
--- a/Utils.jl
+++ b/Utils.jl
@@ -6,19 +6,55 @@ import DSP
 import Printf: @sprintf
 import Crayons as C
 
+const BYELLOW = C.Crayon(foreground=:yellow, bold=true)
+const BRED = C.Crayon(foreground=:light_red, bold=true)
+const RESET = C.Crayon(reset=true)
+
+function get_free_parameters_idx(metadata::Dict)
+  m = metadata["parameters"]
+  return filter(i -> !m[i]["fixed"], eachindex(m))
+end
+
+function reduce_to_free_parameters(metadata::Dict, f::Function, params_init::NamedTuple, lb, ub, q)
+  default_params = Float64[params_init...]
+  free_idx = get_free_parameters_idx(metadata)
+  P_reduced = view(default_params, free_idx)
+  lb_reduced = lb[free_idx]
+  ub_reduced = ub[free_idx]
+  function f_reduced(P::AbstractArray{Float64})
+    P_reduced .= P
+    return f(default_params, q)
+  end
+  return f_reduced, P_reduced, lb_reduced, ub_reduced
+end
+
+function reduce_to_free_parameters(metadata::Dict, params::Vector{Float64})
+  free_idx = get_free_parameters_idx(metadata)
+  return params[free_idx]
+end
+
 function load_init_params(filename::String)
   meta = parsefile(filename)["metadata"]
+  free_idc = get_free_parameters_idx(meta)
   param_names = [p["symbol"] for p in meta["parameters"]]
   param_initial = [p["initial_value"] for p in meta["parameters"]]
   params = NamedTuple{Tuple(map(Symbol, param_names))}(param_initial)
   lower_bounds = [p["min"] for p in meta["parameters"]]
   upper_bounds = [p["max"] for p in meta["parameters"]]
+  free_idx = 1
   println("Got $(length(param_names)) parameters:")
   for (i, k) in enumerate(eachindex(param_names))
     lb_str = @sprintf "%.3g" lower_bounds[i]
     ub_str = @sprintf "%.3g" upper_bounds[i]
     p = @sprintf "%.3g" param_initial[i]
-    println("$(lpad(param_names[i], 7)): $(lpad(lb_str, 7)) <= $(lpad(p, 7)) <= $(lpad(ub_str, 7)) ")
+    if i in free_idc
+      free_idx_str = "$(lpad(free_idx, 2))"
+      free_idx += 1
+    else
+      free_idx_str = "  "
+    end
+
+    println("[", BYELLOW, free_idx_str, RESET, "/$(lpad(i, 2))] $(lpad(param_names[i], 7)): $(lpad(lb_str, 7)) <= $(lpad(p, 7)) <= $(lpad(ub_str, 7)) ")
   end
   return (meta=meta, params=params, lower_bounds=lower_bounds, upper_bounds=upper_bounds)
 end
@@ -96,7 +132,7 @@ function chi2(I_data, I_model, err)
   return sum((I_data .- I_model ./ err) .^ 2)
 end
 
-function add_log_barriers(f::Function, constraints::Vector{Tuple{T,T}}; padding_factor=0.1, mode::Symbol=:outer) where {T<:Real}
+function add_log_barriers(f::Function, constraints::Vector{Tuple{T,T}}; padding_factors=0.1, mode::Symbol=:outer) where {T<:Real}
   n = length(constraints)
 
   @assert mode in [:outer, :inner]
@@ -104,10 +140,10 @@ function add_log_barriers(f::Function, constraints::Vector{Tuple{T,T}}; padding_
   lower_shifts = Float64[]
   upper_shifts = Float64[]
 
-  padding_factors = if padding_factor isa Vector
-    padding_factor
+  pf = if padding_factors isa Vector
+    padding_factors
   else
-    fill(padding_factor, n)
+    fill(padding_factors, n)
   end
 
   k = 0
@@ -120,7 +156,7 @@ function add_log_barriers(f::Function, constraints::Vector{Tuple{T,T}}; padding_
       push!(lower_shifts, lower)
       if isfinite(upper)
         push!(upper_shifts, upper)
-        return padding_factors[k] * (upper - lower)
+        return pf[k] * (upper - lower)
       else
         push!(upper_shifts, Inf)
         return 1.0
@@ -144,6 +180,7 @@ function add_log_barriers(f::Function, constraints::Vector{Tuple{T,T}}; padding_
   end
 
   function barrier(x::Vector{Float64})
+
     lb = (x .- lower_shifts) ./ barrier_widths
     ub = (upper_shifts .- x) ./ barrier_widths
 
@@ -180,14 +217,12 @@ function print_check_bounds(param_names, param_values, lb, ub)
     lb_str = @sprintf "%.3g" lb[i]
     ub_str = @sprintf "%.3g" ub[i]
     p = @sprintf "%.3g" v
-    if lb[i] <= v <= ub[i]
-      color = :white
-      bold = false
+    crayon = if lb[i] <= v <= ub[i]
+      RESET
     else
-      color = :light_red
-      bold = true
+      BRED
     end
-    println("$(lpad(param_names[i], 7)): ", C.Crayon(foreground=color, bold=bold), "$(lpad(lb_str, 7)) <= $(lpad(p, 7)) <= $(lpad(ub_str, 7)) ")
+    println("$(lpad(param_names[i], 7)): ", crayon, "$(lpad(lb_str, 7)) <= $(lpad(p, 7)) <= $(lpad(ub_str, 7)) ")
   end
 end