standardized fracture functions

Author: tzeitim

rogtk/__init__.py

@@ -41,14 +41,14 @@ def nn(expr: IntoExpr, base: int=33) -> pl.Expr:
 #@pl.api.register_expr_namespace("fracture")
 def assemble_sequences(
     expr: IntoExpr,
-    k: int = 31,
-    min_coverage: int = 1,
-    method: str = 'compression',  # Changed from enum to str
-    start_anchor: str | None = None,  # New parameter
-    end_anchor: str | None = None,    # New parameter 
+    k: int = 10,
+    min_coverage: int = 5,
+    method: str = 'shortest_path',  
+    start_anchor: str | None = None,
+    end_anchor: str | None = None, 
+    min_length: int | None = None,
     export_graphs: bool = False,
     only_largest: bool = False,
-    min_length: int | None = None,
     auto_k: bool = False,
     prefix: str | None = None
 ) -> pl.Expr:
@@ -83,9 +83,9 @@ def assemble_sequences(
             "method": method,
             "start_anchor": start_anchor,
             "end_anchor": end_anchor,
+            "min_length": min_length,
             "export_graphs": export_graphs,
             "only_largest": only_largest,
-            "min_length": min_length,
             "auto_k": auto_k,
             "prefix": prefix
         },
@@ -101,11 +101,13 @@ def sweep_assembly_params(
     cov_start: int = 1,
     cov_end: int = 150,
     cov_step: int = 1,
-    method: str = 'compression',  # Changed from enum to str
-    start_anchor: str | None = None,  # New parameter
-    end_anchor: str | None = None,    # New parameter
+    method: str = 'shortest_path',
+    start_anchor: str | None = None,
+    end_anchor: str | None = None,
+    min_length: int | None = None,
     export_graphs: bool = False,
-    prefix: str | None = None
+    prefix: str | None = None,
+    auto_k: bool = False,
 ) -> pl.Expr:
     """
     Run sequence assembly across ranges of k-mer size and minimum coverage parameters.
@@ -135,32 +137,44 @@ def sweep_assembly_params(
             "method": method,
             "start_anchor": start_anchor,
             "end_anchor": end_anchor,
+            "min_length": min_length,
             "export_graphs": export_graphs,
-            "prefix": prefix
+            "prefix": prefix,
+            "auto_k": auto_k
         },
         returns_scalar=True,
         is_elementwise=False,
     )
 
 def optimize_assembly(
     expr: IntoExpr,
-    start_k: int,
-    start_min_coverage: int,
-    start_anchor: str,
-    end_anchor: str,
+    method: str = 'shortest_path',
+    start_anchor: str | None = None,
+    end_anchor: str | None = None,
+    start_k: int = 31,
+    start_min_coverage: int = 1,
+    min_length: int | None = None,
+    export_graphs: bool = False,
+    prefix: str | None = None,
     max_iterations: int | None = None,
     explore_k: bool | None = None,
     prioritize_length: bool | None = None,
 ) -> pl.Expr:
+    if start_anchor is None or end_anchor is None:
+        raise ValueError("Both start_anchor and e®d_anchor are required")
     return register_plugin_function(
         plugin_path=Path(__file__).parent,
         function_name="optimize_assembly_expr",
         args=expr,
         kwargs={
-            "start_k": start_k,
-            "start_min_coverage": start_min_coverage,
+            "method": method,
             "start_anchor": start_anchor,
             "end_anchor": end_anchor,
+            "start_k": start_k,
+            "start_min_coverage": start_min_coverage,
+            "min_length": min_length,
+            "export_graphs": export_graphs,
+            "prefix": prefix,
             "max_iterations": max_iterations,
             "explore_k": explore_k,
             "prioritize_length": prioritize_length

src/expressions.rs

@@ -231,12 +231,18 @@ struct AssemblyKwargs {
     k: usize,
     min_coverage: usize,
     method: String,
+    // optional parameters
     start_anchor: Option<String>,
     end_anchor: Option<String>,
-    export_graphs: Option<bool>,
     min_length: Option<usize>,
+    export_graphs: Option<bool>,
+    only_largest: Option<bool>,
     auto_k: Option<bool>,
     prefix: Option<String>,
+    // optimization related
+    max_iterations: Option<bool>,
+    explore_k: Option<bool>,
+    prioritize_lenth: Option<bool>,
 }
 
 // Default string output type for the expression
@@ -321,8 +327,10 @@ struct SweepParams {
     method: String,
     start_anchor: Option<String>,
     end_anchor: Option<String>,
+    min_length: Option<usize>,
     export_graphs: Option<bool>,
-    prefix: Option<String>
+    prefix: Option<String>,
+    auto_k: Option<bool>,
 }
 
 fn struct_output_type(input_fields: &[Field]) -> PolarsResult<Field> {

src/fracture_opt.rs

@@ -9,15 +9,19 @@ use log::*;
 
 #[derive(Deserialize)]
 pub struct OptimizeParams {
+    pub start_anchor: Option<String>,
+    pub end_anchor: Option<String>,
+
+    pub method: String,
+    pub min_length: Option<usize>,
+    pub export_graphs: Option<bool>,
+    pub prefix: Option<String>,
+
     pub start_k: usize,
     pub start_min_coverage: usize,
-    pub start_anchor: String,
-    pub end_anchor: String,
     pub max_iterations: Option<usize>,
     pub explore_k: Option<bool>,
     pub prioritize_length: Option<bool>,
-    pub method: AssemblyMethod,
-
 }
 
 mod types {
@@ -289,6 +293,18 @@ pub fn optimize_assembly_expr(inputs: &[Series], kwargs: OptimizeParams) -> Pola
     let explore_k = kwargs.explore_k.unwrap_or(false);
     let prioritize_length = kwargs.prioritize_length.unwrap_or(false);
 
+    let start_anchor = kwargs.start_anchor.ok_or_else(||
+        PolarsError::ComputeError("start_anchor is required".into()))?;
+
+    let end_anchor = kwargs.end_anchor.ok_or_else(||
+        PolarsError::ComputeError("end_anchor is required".into()))?;
+
+    let assembly_method = AssemblyMethod::from_str(
+        &kwargs.method,
+        Some(start_anchor.clone()),  
+        Some(end_anchor.clone())
+    ).map_err(|e| PolarsError::ComputeError(e.into()))?;
+
     let start_params = ParamPoint {
         k: kwargs.start_k,
         min_coverage: kwargs.start_min_coverage,
@@ -300,12 +316,12 @@ pub fn optimize_assembly_expr(inputs: &[Series], kwargs: OptimizeParams) -> Pola
     match optimize_assembly(
         &sequences,
         start_params,
-        &kwargs.start_anchor,
-        &kwargs.end_anchor,
+        &start_anchor,
+        &end_anchor,
         max_iterations,
         explore_k,
         prioritize_length,
-        kwargs.method,
+        assembly_method,
     ) {
         Ok(Some(result)) => {
             info!("Optimization successful");