WIP: update to nsa flow 0.7.5

Author: stnava

NAMESPACE

@@ -332,6 +332,7 @@ export(n4BiasFieldCorrection)
 export(neg_violation)
 export(networkEiganat)
 export(nsa_flow)
+export(nsa_flow_autograd)
 export(nsa_flow_fa)
 export(nsa_flow_fa_diagram)
 export(nsa_flow_flowchart)

R/nsa_flow_torch.R

@@ -533,3 +533,228 @@ nsa_flow_torch_ag <- function(
     plot = if (plot) energy_plot else NULL
   )
 }
+
+
+
+#' @title NSA-Flow Optimization via PyTorch AutoGrad (R wrapper)
+#'
+#' @description
+#' Mirror of Python `nsa_flow_autograd()` but callable from R via reticulate.
+#' Uses the Python implementation (nsa_flow.nsa_flow_autograd) and returns
+#' R-friendly results (matrix, data.frame, ggplot).
+#'
+#' @param Y0 numeric matrix p x k initial guess
+#' @param X0 numeric matrix p x k target (or NULL to initialize from Y0)
+#' @param w numeric in [0,1] weighting fidelity vs orthogonality
+#' @param retraction character retraction method
+#' @param max_iter integer max iterations
+#' @param tol numeric convergence tolerance
+#' @param verbose logical
+#' @param seed integer random seed
+#' @param apply_nonneg logical or 'softplus'/'none' etc.
+#' @param optimizer character optimizer name (e.g. 'Adam','lars','sgdp')
+#' @param initial_learning_rate NULL (auto), numeric, or character strategy string
+#' @param lr_strategy character passed to Python if initial_learning_rate is NULL/'auto'
+#' @param fidelity_type character ('basic','scale_invariant','symmetric','normalized')
+#' @param orth_type character ('basic','scale_invariant')
+#' @param record_every integer frequency of recording traces
+#' @param window_size integer window for energy stability
+#' @param plot logical produce ggplot (default FALSE)
+#' @param precision 'float32' or 'float64'
+#' @return list: Y (matrix), traces (data.frame), final_iter, best_total_energy, best_Y_iteration, plot (ggplot or NULL), settings
+#' @export
+nsa_flow_autograd <- function(
+  Y0,
+  X0 = NULL,
+  w = 0.5,
+  retraction = c("soft_polar", "polar", "none"),
+  max_iter = 500L,
+  tol = 1e-6,
+  verbose = FALSE,
+  seed = 42L,
+  apply_nonneg = TRUE,
+  optimizer = "Adam",
+  initial_learning_rate = NULL,
+  lr_strategy = "auto",
+  fidelity_type = "scale_invariant",
+  orth_type = "scale_invariant",
+  record_every = 1L,
+  window_size = 5L,
+  plot = FALSE,
+  precision = "float64"
+) {
+  # basic checks
+  if (!is.matrix(Y0)) stop("Y0 must be a numeric matrix.")
+  if (!is.null(X0) && !is.matrix(X0)) stop("X0 must be a numeric matrix or NULL.")
+  retraction <- match.arg(retraction)
+
+  # import python modules
+  reticulate::py_config() # (optional debug) - remove if noisy
+  torch <- reticulate::import("torch", convert = FALSE)
+  pynsa <- reticulate::import("nsa_flow", convert = FALSE)
+
+  if (is.null(pynsa)) stop("Could not import python package 'nsa_flow'")
+
+  # convert Y0 and X0 to torch tensors (no convert)
+  dtype_str <- ifelse(precision == "float32", "float32", "float64")
+  torch_dtype <- if (dtype_str == "float32") torch$float32 else torch$float64
+
+  Y_torch <- torch$tensor(Y0, dtype = torch_dtype)
+  X_torch <- if (is.null(X0)) reticulate::r_to_py(NULL) else torch$tensor(X0, dtype = torch_dtype)
+
+  # prepare initial_learning_rate argument for Python call
+  if (is.null(initial_learning_rate)) {
+    py_init_lr <- reticulate::r_to_py(NULL)
+    py_lr_strategy <- lr_strategy
+  } else if (is.character(initial_learning_rate)) {
+    # pass through strategy string (e.g. "armijo", "auto")
+    py_init_lr <- reticulate::r_to_py(initial_learning_rate)
+    py_lr_strategy <- initial_learning_rate
+  } else if (is.numeric(initial_learning_rate)) {
+    py_init_lr <- reticulate::r_to_py(as.double(initial_learning_rate))
+    py_lr_strategy <- NA_character_
+  } else {
+    stop("initial_learning_rate must be NULL, a numeric, or a strategy string.")
+  }
+
+  # normalize apply_nonneg options: allow TRUE, FALSE, "softplus", "none"
+  apply_nonneg_py <- switch(
+    as.character(apply_nonneg),
+    "TRUE" = reticulate::r_to_py(TRUE),
+    "FALSE" = reticulate::r_to_py(FALSE),
+    "softplus" = reticulate::r_to_py("softplus"),
+    "none" = reticulate::r_to_py("none"),
+    reticulate::r_to_py(apply_nonneg)
+  )
+
+  # Call python function; be defensive about argument names (match Python signature)
+  py_args <- list(
+    Y0 = Y_torch,
+    X0 = X_torch,
+    w = as.double(w),
+    retraction = retraction,
+    max_iter = as.integer(max_iter),
+    tol = as.double(tol),
+    verbose = reticulate::r_to_py(verbose),
+    seed = as.integer(seed),
+    apply_nonneg = apply_nonneg_py,
+    optimizer = optimizer,
+    initial_learning_rate = py_init_lr,
+    lr_strategy = lr_strategy,
+    fidelity_type = fidelity_type,
+    orth_type = orth_type,
+    record_every = as.integer(record_every),
+    window_size = as.integer(window_size),
+    precision = precision
+  )
+
+  # Try to call python function and handle errors clearly
+  res_py <- tryCatch(
+    {
+      do.call(pynsa$nsa_flow_autograd, py_args)
+    },
+    error = function(e) {
+      stop("Error calling Python nsa_flow_autograd():\n", e$message)
+    }
+  )
+
+  # Convert outputs
+  # Expect res_py to be a dict-like object with keys "Y", "traces", "final_iter", "best_total_energy", "best_Y_iteration", "target", "settings"
+  # Convert Y (torch tensor) -> numeric matrix
+  Y_out <- NULL
+  if (!is.null(res_py$Y)) {
+    # res_py$Y might be a torch tensor; convert safely
+    # Use detach() if available
+    try({
+      # If it's a tensor object, call detach().numpy(); else, try py_to_r
+      if (!is.null(res_py$Y$detach)) {
+        Y_out <- as.matrix(res_py$Y$detach()$cpu()$numpy())
+      } else {
+        Y_out <- reticulate::py_to_r(res_py$Y)
+      }
+    }, silent = TRUE)
+
+    # fallback
+    if (is.null(Y_out)) {
+      Y_out <- tryCatch(reticulate::py_to_r(res_py$Y), error = function(e) NULL)
+    }
+  }
+
+  # Convert traces to data.frame if possible
+  traces_df <- NULL
+  if (!is.null(res_py$traces)) {
+    traces_df <- tryCatch({
+      reticulate::py_to_r(res_py$traces)
+    }, error = function(e) {
+      # If traces is a list of dicts, convert manually
+      tlist <- reticulate::py_to_r(res_py$traces)
+      if (is.list(tlist) && length(tlist) > 0 && is.list(tlist[[1]])) {
+        do.call(rbind, lapply(tlist, function(x) as.data.frame(x, stringsAsFactors = FALSE)))
+      } else {
+        as.data.frame(tlist)
+      }
+    })
+    # ensure rownames removed
+    rownames(traces_df) <- NULL
+  }
+
+  # final iter & energy
+  final_iter <- tryCatch(reticulate::py_to_r(res_py$final_iter), error = function(e) NA)
+  best_total_energy <- tryCatch(reticulate::py_to_r(res_py$best_total_energy), error = function(e) NA)
+  best_Y_iteration <- tryCatch(reticulate::py_to_r(res_py$best_Y_iteration), error = function(e) NA)
+  settings <- tryCatch(reticulate::py_to_r(res_py$settings), error = function(e) NULL)
+
+  # Build ggplot trace if requested and data available
+  energy_plot <- NULL
+  if (plot && !is.null(traces_df) && nrow(traces_df) > 0) {
+    if (!("fidelity" %in% colnames(traces_df)) || !("orthogonality" %in% colnames(traces_df))) {
+      # try to coerce likely-named columns
+      possible_fid <- grep("fid", names(traces_df), value = TRUE, ignore.case = TRUE)
+      possible_orth <- grep("orth", names(traces_df), value = TRUE, ignore.case = TRUE)
+      if (length(possible_fid) >= 1) names(traces_df)[which(names(traces_df) == possible_fid[1])] <- "fidelity"
+      if (length(possible_orth) >= 1) names(traces_df)[which(names(traces_df) == possible_orth[1])] <- "orthogonality"
+    }
+
+    if ("fidelity" %in% colnames(traces_df) && "orthogonality" %in% colnames(traces_df)) {
+      max_fid <- max(traces_df$fidelity, na.rm = TRUE)
+      max_orth <- max(traces_df$orthogonality, na.rm = TRUE)
+      ratio <- if (max_orth > 0) max_fid / max_orth else 1
+
+      energy_plot <- ggplot2::ggplot(traces_df, ggplot2::aes(x = iter)) +
+        ggplot2::geom_line(ggplot2::aes(y = fidelity, color = "Fidelity"), size = 1.1) +
+        ggplot2::geom_point(ggplot2::aes(y = fidelity, color = "Fidelity"), size = 1.2, alpha = 0.7) +
+        ggplot2::geom_line(ggplot2::aes(y = orthogonality * ratio, color = "Orthogonality"), size = 1.1) +
+        ggplot2::geom_point(ggplot2::aes(y = orthogonality * ratio, color = "Orthogonality"), size = 1.2, alpha = 0.7) +
+        ggplot2::scale_y_continuous(name = "Fidelity Energy",
+                                    sec.axis = ggplot2::sec_axis(~ . / ratio, name = "Orthogonality Defect")) +
+        ggplot2::scale_color_manual(values = c("Fidelity" = "#1f78b4", "Orthogonality" = "#33a02c")) +
+        ggplot2::labs(title = paste("NSA-Flow Optimization Trace:", retraction),
+                      subtitle = paste0("fidelity_type=", fidelity_type, ", orth_type=", orth_type),
+                      x = "Iteration", color = "Term") +
+        ggplot2::theme_minimal(base_size = 13) +
+        ggplot2::theme(plot.title = ggplot2::element_text(face = "bold", hjust = 0.5),
+                       legend.position = "top")
+    }
+  }
+
+  # Rescale target/back to original magnitude if Python returned scaled Y or target info:
+  # If Python returned 'target' or 'settings' containing scale_ref, you could rescale;
+  # here we assume outputs are in original scale or the Python function already rescaled.
+  # Provide Y_out as numeric matrix; maintain dimnames from input
+  if (!is.null(Y_out)) {
+    dimnames(Y_out) <- list(rownames(Y0), colnames(Y0))
+  }
+
+  out <- list(
+    Y = Y_out,
+    traces = traces_df,
+    final_iter = final_iter,
+    best_total_energy = best_total_energy,
+    best_Y_iteration = best_Y_iteration,
+    plot = if (plot) energy_plot else NULL,
+    settings = settings
+  )
+
+  class(out) <- c("nsa_flow_result", class(out))
+  return(out)
+}

docs/compare_optimizers_for_NSA-Flow.Rmd

@@ -27,7 +27,7 @@ if (exists("matoption") && matoption == "wide") mat1 <- t(mat1)
 ws <- seq(0.05, 0.95, by = 0.05)
 optimizers <- list_simlr_optimizers( torch=TRUE )
 goptimizers = optimizers[ !(optimizers %in% c("adamp","lbfgs","padam","sgdp","adabound","adamax"))]
-max_iter <- 500
+max_iter <- 50
 tol <- 1e-6
 def_ret <- "soft_polar"
 if (!exists("raw_results")) {
@@ -37,7 +37,7 @@ if (!exists("raw_results")) {
 #      print(paste("Running optimizer:", o, "with w =", round(w, 2)))
       start_time <- Sys.time()
       res = nsa_flow_torch_ag(mat1, w = w, verbose = F, retraction = def_ret,
-          optimizer = o, initial_learning_rate=NA )
+          optimizer = o, initial_learning_rate=NA, max_iter=max_iter, tol=tol)
       end_time <- Sys.time()
       runtime <- as.numeric(difftime(end_time, start_time, units = "secs"))
       if (!is.null(res)) {