Spaces:
Running
Running
File size: 4,716 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 |
"""
U1FieldNode (Electromagnetism Metaphor)
Simulates a U(1) gauge force, like electromagnetism.
It takes a grayscale "charge density" map and calculates
the resulting force field (like an E-field).
[FIXED] Initialized self.potential in __init__ and
saved potential to self.potential in step().
"""
import numpy as np
import cv2
# --- Magic import block ---
import __main__
BaseNode = __main__.BaseNode
QtGui = __main__.QtGui
# --------------------------
class U1FieldNode(BaseNode):
"""
Generates a U(1) force field from a charge density map.
"""
NODE_CATEGORY = "Generator"
NODE_COLOR = QtGui.QColor(100, 150, 220) # Blue
def __init__(self, size=128):
super().__init__()
self.node_title = "U(1) Field (E/M)"
self.inputs = {
'charge_in': 'image', # Grayscale image (0-1)
'strength': 'signal' # 0-1, force strength
}
self.outputs = {
'potential_out': 'image', # The scalar potential (blurred charge)
'field_viz': 'image' # Vector field visualization
}
self.size = int(size)
# --- START FIX ---
# Initialize the output variables to prevent AttributeError
self.viz = np.zeros((self.size, self.size, 3), dtype=np.float32)
self.potential = np.zeros((self.size, self.size), dtype=np.float32)
# --- END FIX ---
def _prepare_image(self, img):
if img is None:
return np.full((self.size, self.size), 0.5, dtype=np.float32)
if img.dtype != np.float32: img = img.astype(np.float32)
if img.max() > 1.0: img /= 255.0
img_resized = cv2.resize(img, (self.size, self.size),
interpolation=cv2.INTER_LINEAR)
if img_resized.ndim == 3:
return cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
return img_resized
def step(self):
# --- 1. Get Charge Density ---
# Map input [0, 1] to charge [-1, 1]
charge_density = (self._prepare_image(
self.get_blended_input('charge_in', 'first')
) * 2.0) - 1.0
strength = self.get_blended_input('strength', 'sum') or 1.0
# --- 2. Calculate Potential ---
# Simulate long-range 1/r potential by blurring
# A large blur kernel simulates the 1/r falloff
ksize = self.size // 4 * 2 + 1 # Must be odd
self.potential = cv2.GaussianBlur(charge_density, (ksize, ksize), 0)
# --- 3. Calculate Force Field (E-Field) ---
# E = -∇V (Force is the negative gradient of potential)
grad_x = -cv2.Sobel(self.potential, cv2.CV_32F, 1, 0, ksize=3) * strength
grad_y = -cv2.Sobel(self.potential, cv2.CV_32F, 0, 1, ksize=3) * strength
# --- 4. Create Visualization ---
self.viz = np.zeros((self.size, self.size, 3), dtype=np.float32)
step = 8 # Draw an arrow every 8 pixels
for y in range(0, self.size, step):
for x in range(0, self.size, step):
vx = grad_x[y, x] * 20 # Scale for viz
vy = grad_y[y, x] * 20
pt1 = (x, y)
pt2 = (int(np.clip(x + vx, 0, self.size-1)),
int(np.clip(y + vy, 0, self.size-1)))
# Color based on direction
angle = np.arctan2(vy, vx) + np.pi
hue = int(angle / (2 * np.pi) * 179) # 0-179 for OpenCV HSV
color_hsv = np.uint8([[[hue, 255, 255]]])
color_rgb = cv2.cvtColor(color_hsv, cv2.COLOR_HSV2RGB)[0][0]
color_float = color_rgb.astype(np.float32) / 255.0
# --- START FIX ---
# Convert numpy.float32 to standard Python floats for OpenCV
color_tuple = (float(color_float[0]), float(color_float[1]), float(color_float[2]))
cv2.arrowedLine(self.viz, pt1, pt2, color_tuple, 1, cv2.LINE_AA)
# --- END FIX ---
def get_output(self, port_name):
if port_name == 'potential_out':
# Normalize potential [-max, +max] to [0, 1]
p_max = np.max(np.abs(self.potential))
if p_max == 0: return np.full((self.size, self.size), 0.5, dtype=np.float32)
return (self.potential / (2 * p_max)) + 0.5
elif port_name == 'field_viz':
return self.viz
return None
def get_display_image(self):
return self.viz |