continual-foragax 0.40.0__py3-none-any.whl → 0.41.0__py3-none-any.whl

foragax/objects.py

@@ -327,15 +327,18 @@ class FourierObject(BaseForagaxObject):
 
         # Extract interleaved coefficients: [a1, b1, a2, b2, ...]
         ab_coeffs = params[3:]
+        # Reuse existing arrays via reshaping/slices for vectorization
+        # Reshape ab_coeffs to (num_terms, 2) where col 0 is a_n, col 1 is b_n
         n_terms = len(ab_coeffs) // 2
+        coeffs = ab_coeffs.reshape((n_terms, 2))
+        a_coeffs = coeffs[:, 0]
+        b_coeffs = coeffs[:, 1]
 
-        # Compute Fourier series: sum(a_n*cos(n*t) + b_n*sin(n*t))
-        reward = 0.0
-        for i in range(n_terms):
-            freq = i + 1
-            a_i = ab_coeffs[2 * i]  # a coefficient at index 2i
-            b_i = ab_coeffs[2 * i + 1]  # b coefficient at index 2i+1
-            reward += a_i * jnp.cos(freq * t) + b_i * jnp.sin(freq * t)
+        freqs = jnp.arange(1, n_terms + 1, dtype=jnp.float32)
+        terms = freqs * t
+
+        # Calculate sum(a_n * cos(n*t) + b_n * sin(n*t))
+        reward = jnp.sum(a_coeffs * jnp.cos(terms) + b_coeffs * jnp.sin(terms))
 
         # Apply min-max normalization to [-1, 1], then scale by base_magnitude
         # Formula: 2 * (x - min) / (max - min) - 1

foragax/rendering.py

@@ -38,39 +38,15 @@ def apply_true_borders(
         jax.image.ResizeMethod.NEAREST,
     )
 
-    # Create border mask (2-pixel thick borders) - vectorized like grid lines
-    height, width = grid_size
-    img_height, img_width = height * 24, width * 24
-
-    border_mask = jnp.zeros((img_height, img_width), dtype=bool)
-
-    # Create border row and column indices for all cells at once
-    cell_rows = jnp.arange(height)
-    cell_cols = jnp.arange(width)
-
-    # Top border rows: 2 rows per cell
-    top_border_rows = cell_rows[:, None] * 24 + jnp.arange(2)[None, :]
-    top_border_rows_flat = top_border_rows.flatten()
-
-    # Bottom border rows: 2 rows per cell
-    bottom_border_rows = cell_rows[:, None] * 24 + 22 + jnp.arange(2)[None, :]
-    bottom_border_rows_flat = bottom_border_rows.flatten()
-
-    # Left border columns: 2 columns per cell
-    left_border_cols = cell_cols[:, None] * 24 + jnp.arange(2)[None, :]
-    left_border_cols_flat = left_border_cols.flatten()
-
-    # Right border columns: 2 columns per cell
-    right_border_cols = cell_cols[:, None] * 24 + 22 + jnp.arange(2)[None, :]
-    right_border_cols_flat = right_border_cols.flatten()
-
-    # Set top and bottom borders (full width rectangles)
-    all_border_rows = jnp.concatenate([top_border_rows_flat, bottom_border_rows_flat])
-    border_mask = border_mask.at[all_border_rows, :].set(True)
-
-    # Set left and right borders (full height rectangles)
-    all_border_cols = jnp.concatenate([left_border_cols_flat, right_border_cols_flat])
-    border_mask = border_mask.at[:, all_border_cols].set(True)
+    # Create border mask (2-pixel thick borders) using vectorized modulo operations
+    img_height, img_width = grid_size[0] * 24, grid_size[1] * 24
+    y_idx = jnp.arange(img_height) % 24
+    x_idx = jnp.arange(img_width) % 24
+
+    # Border pixels are those with offset 0, 1, 22, or 23 within each 24x24 cell
+    is_border_row = (y_idx < 2) | (y_idx >= 22)
+    is_border_col = (x_idx < 2) | (x_idx >= 22)
+    border_mask = is_border_row[:, None] | is_border_col[None, :]
 
     # Apply border mask: use HSV border colors for border pixels, base colors elsewhere
     result_img = jnp.where(border_mask[..., None], border_img, base_img)