Implemented the float format in cast_coef_to_string() instead.

Bravos-Power · Apr 10, 2024 · dcb2dac · dcb2dac
1 parent a74b83e
commit dcb2dac
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 12 deletions.
diff --git a/src/pyoframe/constraints.py b/src/pyoframe/constraints.py
@@ -577,15 +577,11 @@ def to_str_table(
         data = self.data if include_const_term else self.variable_terms
         if var_map is None:
             var_map = self._model.var_map if self._model is not None else DEFAULT_MAP
-        data = cast_coef_to_string(data)
+        data = cast_coef_to_string(data, float_format=float_format)
         data = var_map.map_vars(data)
         dimensions = self.dimensions
 
         # Create a string for each term
-        if float_format is not None:
-            data = data.with_columns(
-            pl.col(COEF_KEY).map_batches(lambda x: sprintf(x, float_format), return_dtype=str)
-            )
 
         data = (
             data
@@ -666,10 +662,9 @@ def to_str(
         [1]: 3.141592653589793 x1
         [2]: 3.141592653589793 x2
         [3]: 3.141592653589793 x3
-        >>> print((x*np.pi >= 0).to_str(float_format="%0.4f"))
-        [1]: 3.1416 x1 >= 0
-        [2]: 3.1416 x2 >= 0
-        [3]: 3.1416 x3 >= 0
+        >>> print((np.pi*x.drop_unmatched() -  2*x.next("t") >= 0).to_str(float_format="%0.4f"))
+        [1]: 3.1416 x1 -2 x2 >= 0
+        [2]: 3.1416 x2 -2 x3 >= 0
         """
         result = ""
         if include_header:

diff --git a/src/pyoframe/util.py b/src/pyoframe/util.py
@@ -156,7 +156,7 @@ def concat_dimensions(
 
 
 def cast_coef_to_string(
-    df: pl.DataFrame, column_name: str = COEF_KEY, drop_ones=True
+    df: pl.DataFrame, column_name: str = COEF_KEY, drop_ones=True, float_format=None
 ) -> pl.DataFrame:
     """
     Examples:
@@ -183,7 +183,8 @@ def cast_coef_to_string(
     df = df.with_columns(
         pl.when(pl.col(column_name) == pl.col(column_name).floor())
         .then(pl.col(column_name).cast(pl.Int64).cast(pl.String))
-        .otherwise(pl.col(column_name).cast(pl.String))
+        #.otherwise(pl.col(column_name).cast(pl.String))
+        .otherwise(pl.col(column_name).map_batches(lambda x: sprintf(x, float_format), str))
         .alias(column_name)
     )
 
@@ -203,9 +204,22 @@ def cast_coef_to_string(
         "_sign"
     )
 
+def to_string_with_precision(s: pl.Series, float_precision: int):
+    """
+    Examples
+    --------
+    >>> import polars as pl
+    >>> s = pl.Series([1, 2, 3.4, 5.6789])
+    >>> print(to_string_with_precision(s,2).to_list())
+    ['1.00', '2.00', '3.40', '5.68']
+
+    """
+    return sprintf(s, f"%0.{float_precision}f")
+
+
 
 # Function by mcrumiller (https://github.com/mcrumiller). See https://github.com/pola-rs/polars/issues/7133
-def sprintf(s, fmt):
+def sprintf(s: pl.Series, fmt: str):
     """
     Formats each element of a Polars Series `s` according to the format specifier `fmt`,
     similarly to sprintf in C or format in Python. It supports basic formatting for
@@ -244,6 +258,9 @@ def sprintf(s, fmt):
     ['1.00', '2.00', '3.40', '5.68']
     """
 
+    if fmt is None:
+        return s.cast(pl.String)
+
     # parse format
     parser = re.compile(r"^%(?P<pct>%?)(?P<align>[\<\>|]?)(?P<head>\d*)(?P<dot>\.?)(?P<dec>\d*)(?P<char>[dfs])$")
     result = parser.match(fmt)