File size: 7,105 Bytes
3bb804c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
"""

Antti's PhiWorld Node - A TADS-like particle field simulation

Driven by an energy signal and perturbed by an image.

Based on the physics from phiworld2.py.

Requires: pip install scipy

Place this file in the 'nodes' folder

"""

import numpy as np
from PyQt6 import QtGui
import cv2

import sys
import os
# --- This is the new, correct block ---
import __main__
BaseNode = __main__.BaseNode
PA_INSTANCE = getattr(__main__, "PA_INSTANCE", None)
# ------------------------------------

try:
    from scipy.signal import convolve2d
    from scipy.ndimage import maximum_filter
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False
    print("Warning: PhiWorldNode requires 'scipy'.")
    print("Please run: pip install scipy")

class PhiWorldNode(BaseNode):
    NODE_CATEGORY = "Transform"
    NODE_COLOR = QtGui.QColor(220, 120, 40) # Transform Orange
    
    def __init__(self, grid_size=96):
        super().__init__()
        self.node_title = "Antti's PhiWorld"
        
        self.inputs = {
            'energy_in': 'signal', # Drives the simulation
            'perturb_in': 'image'  # Pushes the field
        }
        self.outputs = {
            'field': 'image',       # The raw phi field
            'particles': 'image',   # Just the detected particles
            'count': 'signal'       # Number of particles
        }
        
        self.grid_size = int(grid_size)
        
        # --- Parameters from phiworld2.py ---
        self.dt = 0.08
        self.damping = 0.005 # Increased damping for stability in node
        self.base_c_sq = 1.0
        self.tension_factor = 5.0
        self.potential_lin = 1.0
        self.potential_cub = 0.2
        self.biharmonic_gamma = 0.02
        self.particle_threshold = 0.5
        
        # --- Internal State ---
        self.phi = np.zeros((self.grid_size, self.grid_size), dtype=np.float64)
        self.phi_old = np.zeros_like(self.phi)
        
        # Optimized Laplacian Kernel
        self.laplacian_kernel = np.array([[0, 1, 0], [1, -4, 1]], dtype=np.float64)
        
        # Outputs
        self.particle_image = np.zeros_like(self.phi, dtype=np.float32)
        self.particle_count = 0.0

        if not SCIPY_AVAILABLE:
            self.node_title = "PhiWorld (No SciPy!)"

    # --- Physics methods adapted from phiworld2.py ---
    
    def _laplacian(self, f):
        # Using np.roll is faster than convolve2d for this kernel
        lap_x = np.roll(f, -1, axis=1) - 2 * f + np.roll(f, 1, axis=1)
        lap_y = np.roll(f, -1, axis=0) - 2 * f + np.roll(f, 1, axis=0)
        return lap_x + lap_y

    def _biharmonic(self, f):
        lap_f = self._laplacian(f)
        return self._laplacian(lap_f) # Laplacian of the Laplacian

    def _potential_deriv(self, phi):
        return (-self.potential_lin * phi
                + self.potential_cub * (phi**3))

    def _local_speed_sq(self, phi):
        intensity = phi**2
        return self.base_c_sq / (1.0 + self.tension_factor * intensity + 1e-9)

    def _track_particles(self, field):
        """Optimized particle tracking using scipy.ndimage."""
        # Find local maxima using a 3x3 filter
        maxima_mask = (field == maximum_filter(field, size=3))
        # Find points above threshold
        threshold_mask = (field > self.particle_threshold)
        
        # Combine masks
        particle_mask = (maxima_mask & threshold_mask)
        
        # Update outputs
        self.particle_image = particle_mask.astype(np.float32)
        self.particle_count = np.sum(particle_mask)

    def step(self):
        if not SCIPY_AVAILABLE:
            return

        # Get inputs
        energy = self.get_blended_input('energy_in', 'sum') or 0.0
        perturb_img = self.get_blended_input('perturb_in', 'mean')
        
        if energy <= 0.01:
            # If no energy, dampen the field
            self.phi *= (1.0 - (self.damping * 10)) # Faster damping
            self.phi_old = self.phi.copy()
            self.particle_image *= 0.9
            self.particle_count = 0
            return

        # --- Run simulation step (from phiworld2.py) ---
        
        # Calculate forces
        lap_phi = self._laplacian(self.phi)
        biharm_phi = self._biharmonic(self.phi)
        c2 = self._local_speed_sq(self.phi)
        V_prime = self._potential_deriv(self.phi)
        
        # Scale acceleration by energy input
        acceleration = energy * ( (c2 * lap_phi) - V_prime - (self.biharmonic_gamma * biharm_phi) )

        # Update field (Verlet integration)
        velocity = self.phi - self.phi_old
        phi_new = self.phi + (1.0 - self.damping * self.dt) * velocity + (self.dt**2) * acceleration

        # --- Add Image Perturbation ---
        if perturb_img is not None:
            # Resize image to grid
            img_resized = cv2.resize(perturb_img, (self.grid_size, self.grid_size),
                                     interpolation=cv2.INTER_AREA)
            # "Push" the field with the image, scaled by energy
            phi_new += (img_resized - 0.5) * 0.1 * energy # (Image is 0-1, so map to -0.5 to 0.5)

        self.phi_old = self.phi.copy()
        self.phi = phi_new
        
        # Clamp to prevent instability
        self.phi = np.clip(self.phi, -10.0, 10.0)

        # Track particles on the new field
        self._track_particles(np.abs(self.phi))

    def get_output(self, port_name):
        if port_name == 'field':
            # Normalize field for output [-2, 2] -> [0, 1]
            return np.clip(self.phi * 0.25 + 0.5, 0.0, 1.0)
        elif port_name == 'particles':
            return self.particle_image
        elif port_name == 'count':
            return self.particle_count
        return None
        
    def get_display_image(self):
        # Normalize field for display
        img_norm = np.clip(self.phi * 0.25 + 0.5, 0.0, 1.0)
        
        img_u8 = (img_norm * 255).astype(np.uint8)
        
        # Apply a colormap
        img_color = cv2.applyColorMap(img_u8, cv2.COLORMAP_VIRIDIS)
        
        # Overlay particles in bright red
        img_color[self.particle_image > 0] = (0, 0, 255) # BGR for red
        
        img_color = np.ascontiguousarray(img_color)
        h, w = img_color.shape[:2]
        return QtGui.QImage(img_color.data, w, h, 3*w, QtGui.QImage.Format.Format_BGR888)

    def get_config_options(self):
        return [
            ("Grid Size", "grid_size", self.grid_size, None),
            ("Particle Thresh", "particle_threshold", self.particle_threshold, None),
            ("Damping", "damping", self.damping, None),
            ("Tension", "tension_factor", self.tension_factor, None),
            ("Linear Pot.", "potential_lin", self.potential_lin, None),
            ("Cubic Pot.", "potential_cub", self.potential_cub, None),
            ("Biharmonic (g)", "biharmonic_gamma", self.biharmonic_gamma, None),
        ]