Spaces:
Running
Running
File size: 2,805 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 |
import __main__
BaseNode = __main__.BaseNode
QtGui = __main__.QtGui
import numpy as np
import cv2
import time
class SpringNode(BaseNode):
"""
Simulates a 1D damped spring.
F = -k*x - c*v
"""
NODE_CATEGORY = "AI / Physics"
NODE_COLOR = QtGui.QColor(100, 180, 100)
def __init__(self, mass=1.0, stiffness=0.1, damping=0.05):
super().__init__()
self.node_title = "Spring (1D)"
self.inputs = {'target_pos': 'signal'}
self.outputs = {'position': 'signal'}
self.mass = float(mass)
self.stiffness = float(stiffness)
self.damping = float(damping)
self.position = 0.0
self.velocity = 0.0
self.last_time = time.time()
def step(self):
# Calculate delta time
current_time = time.time()
dt = current_time - self.last_time
if dt > 0.1: # Clamp large timesteps (e.g., on load)
dt = 0.1
self.last_time = current_time
# Get target
target = self.get_blended_input('target_pos', 'sum') or 0.0
# Calculate forces
displacement = self.position - target
spring_force = -self.stiffness * displacement
damping_force = -self.damping * self.velocity
total_force = spring_force + damping_force
# Update physics (Euler integration)
acceleration = total_force / self.mass
self.velocity += acceleration * dt
self.position += self.velocity * dt
def get_output(self, port_name):
if port_name == 'position':
return self.position
return None
def get_display_image(self):
w, h = 256, 128
img = np.zeros((h, w, 3), dtype=np.uint8)
# Draw position
pos_x = int(np.clip((self.position + 2) / 4.0, 0, 1) * w)
cv2.circle(img, (pos_x, h//2), 10, (0, 255, 0), -1)
# Draw target
target = self.get_blended_input('target_pos', 'sum') or 0.0
target_x = int(np.clip((target + 2) / 4.0, 0, 1) * w)
cv2.circle(img, (target_x, h//2), 5, (0, 0, 255), -1)
cv2.putText(img, f"Pos: {self.position:.2f}", (5, 15),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255,255,255), 1)
cv2.putText(img, f"Vel: {self.velocity:.2f}", (5, 30),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255,255,255), 1)
return QtGui.QImage(img.data, w, h, w*3, QtGui.QImage.Format.Format_RGB888)
def get_config_options(self):
return [
("Mass", "mass", self.mass, None),
("Stiffness", "stiffness", self.stiffness, None),
("Damping", "damping", self.damping, None)
] |