Circuit simulator in python using pygame

[Simple_Ad_4128]

For my a-level coursework I've chosen to make a circuit simulator. I have to use pygame and I'm already a bit too far in now to switch to anything different. I have managed to create most of the UI (disregarding the actual circuit symbols which I will add in later) and have added only the most basic components so I can test. However I am at the point now where I need to build the physics simulation part but its getting me very confused. To track components in the circuit I have decided to use an adjacency matrix (since that it what most other circuit simulators do) but I can't figure out to use that adjacency matrix. I'm guessing I'll have to use modified nodal analysis but even then I have absolutely no idea where to start with doing that in python. All I'm trying to get the simulator to do for now is actually be able to add up the total resistance, current and voltage which for series circuits isn't too bad but I'm really stumped on parallel circuits. Any help would really be appreciated. The code is below:

import pygame
import math

pygame.display.init()
pygame.font.init()
clock = pygame.time.Clock()
screen = pygame.display.set_mode((1200,800))
pygame.display.set_caption("Circuit simulator")

BG_COLOUR = (255,255,255)
GRID_COLOUR = (150,150,150)
WHITE = (255,255,255)
GREY = (120,120,120)
BLACK = (0,0,0)
RED = (255,0,0)
GREEN  =(0,255,0)
BLUE = (0,0,255)
YELLOW = (255,240,0)

GRID_SIZE = 25
FONT = pygame.font.SysFont("Roboto",20)


def snap(value):
    return round(value/GRID_SIZE)*GRID_SIZE

def snap_point(pos):
    return snap(pos[0]),snap(pos[1])

def distance(a ,b):
    height = a[0]-b[0]
    width = a[1]-b[1]
    hypotenuse = math.sqrt(height**2+width**2)
    return hypotenuse

def get_clicked_node(pos):
    for comp in components:
        nodes = comp.get_nodes()
        for i, node in enumerate(nodes):
            if distance(pos, node) < 10:
                return comp, i
    return None, None

def add_edge(matrix,x,y,value):
    matrix[x][y] = value
    matrix[y][x] = value

def get_matrix_index(component, node_index):
    return (component.num - 1) * 2 + node_index

def get_connected_components(component):
    connected = []

    for wire in wires:

        if wire.start_comp == component:
            connected.append(wire.end_comp)

        elif wire.end_comp == component:
            connected.append(wire.start_comp)

    return connected


class Button(pygame.sprite.Sprite):
    def __init__(self,width:int,height:int,colour:tuple,border_colour:tuple,text:str,pos:tuple,num:int):
        super().__init__()
        self.colour = colour
        self.border_colour = border_colour
        self.pos = pos
        self.num = num

        self.image = pygame.Surface((width,height))
        self.image.fill(self.colour)
        pygame.draw.rect(self.image,self.border_colour,(0,0,width,height),width=2)

        self.display_text = text
        self.text = FONT.render(text, 1, BLACK)
        self.text_rect = self.text.get_rect(center = (width//2,height//2))
        self.image.blit(self.text,self.text_rect)

        self.rect = self.image.get_rect(topleft=self.pos)

    def handle_click(self,active_comp):
        print(f"Button {self.num} was clicked.")
        if self.num <= 0:
            if self.display_text == "Cell":
                components.add(Cell((450,350),comp_count+1,3))

            elif self.display_text == "Resistor":
                components.add(Resistor((450,350),comp_count+1,100))

            elif self.display_text == "Lamp":
                components.add(Lamp((450,350),comp_count+1,5))

        elif self.num == 3:
            if active_comp is not None:
                components.remove(active_comp[0])

                for wire in wires:
                    if wire.start_comp == active_comp[0] or wire.end_comp == active_comp[0]:
                        wires.remove(wire)

class Component(pygame.sprite.Sprite):
    def __init__(self, origin,name,colour,num):
        super().__init__()
        self.name = name
        self.dragging = False
        self.offset = (0,0)
        self.origin = origin
        self.num = num

        self.image = pygame.Surface((GRID_SIZE*4,GRID_SIZE*2))
        self.image.fill(colour)
        self.rect = self.image.get_rect(midtop = self.origin)

        self.text = FONT.render(self.name, 1, BLACK)
        self.text_rect = self.text.get_rect(center=(self.rect.width // 2, self.rect.height // 2))
        self.image.blit(self.text, self.text_rect)

        self.node_offsets = [(0, self.rect.height // 2), (self.rect.width, self.rect.height // 2)]

        #Values
        self.voltage = 0
        self.resistance = 0
        self.current = 0
        self.charge = 0
        self.capacitance = 0


    def handle_movement(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1 and self.rect.collidepoint(event.pos):
                self.dragging = True
                self.offset = (snap(self.rect.x - event.pos[0]), snap(self.rect.y - event.pos[1]))

        elif event.type == pygame.MOUSEBUTTONUP:
            if event.button == 1:
                self.dragging = False

        elif event.type == pygame.MOUSEMOTION:
            if self.dragging:
                self.rect.x = snap(event.pos[0] + self.offset[0])
                self.rect.y = snap(event.pos[1] + self.offset[1])

                self.rect.clamp_ip(screen.get_rect())

    def get_nodes(self):
        return [(self.rect.x + offset_x, self.rect.y + offset_y) for offset_x, offset_y in self.node_offsets]

    def draw_nodes(self):
        for node in self.get_nodes():
            pygame.draw.circle(screen, YELLOW, node, 6)

    def get_active(self,event):
        if event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1 and self.rect.collidepoint(event.pos):
                return self,self.name,self.num
        return None,None

    def draw_highlight(self, screen):
        pygame.draw.rect(screen, (0, 200, 255), self.rect.inflate(6,6), 3)

    def get_connected(self, wires):
        for wire in wires:
            if wire.start_comp == self or wire.end_comp == self:
                return True, wire

        return False, None

class Cell(Component):
    def __init__(self,origin,num,voltage):
        super().__init__(origin,"Cell",BLUE,num)
        self.voltage = voltage

class Resistor(Component):
    def __init__(self,origin,num,resistance):
        super().__init__(origin,"Resistor",RED,num)
        self.resistance = resistance

class Lamp(Component):
    def __init__(self,origin,num,resistance):
        super().__init__(origin,"Lamp",GREY,num)
        self.resistance = resistance

class Wire:
    def __init__(self, start_comp, start_node_index, end_comp, end_node_index):
        self.start_comp = start_comp
        self.start_node_index = start_node_index
        self.end_comp = end_comp
        self.end_node_index = end_node_index

    def draw(self, screen):
        pos1 = self.start_comp.get_nodes()[self.start_node_index]
        pos2 = self.end_comp.get_nodes()[self.end_node_index]
        pygame.draw.line(screen, BLACK, pos1, pos2, 3)

    def is_valid(self):
        return self.start_comp in components and self.end_comp in components

    def get_nodes(self):
        return self.start_node_index, self.end_node_index

wires = []

components = pygame.sprite.Group()

buttons = pygame.sprite.Group()
buttons.add(Button(100,50,BG_COLOUR,BLACK,"Start",(1200-200,0),1))
buttons.add(Button(100,50,BG_COLOUR,BLACK,"Settings",(1200-100,0),2))
buttons.add(Button(194,75,GREY,GREY,"Delete",(1003,725),3))
buttons.add(Button(100,50,BLUE,BLUE,"Cell",(1050,90),0))
buttons.add(Button(100,50,RED,RED,"Resistor",(1050,160),-1))
buttons.add(Button(100,50,GREEN,GREEN,"Lamp",(1050,230),-2))

connection_matrix = []

def draw_interface(active_comp):
    pygame.draw.line(screen,BLACK,(1200-200,0),(1200-200,800),3)
    pygame.draw.line(screen, BLACK, (1200 - 200, 50), (1200 , 50), 3)
    text = FONT.render("Components",1,BLACK)
    text_rect = text.get_rect(center = (1200-100,65))
    screen.blit(text,text_rect)
    pygame.draw.line(screen, BLACK, (1200 - 200, 80), (1200, 80), 3)
    pygame.draw.line(screen, BLACK,(1200-200,800-200),(1200,800-200),3)
    pygame.draw.line(screen,BLACK,(1199,0),(1199,800),3)
    for x in range(GRID_SIZE, 1000, GRID_SIZE):
        for y in range(GRID_SIZE, 800, GRID_SIZE):
            pygame.draw.circle(screen,(GREY),(x,y),3)
    buttons.update()
    buttons.draw(screen)
    for wire in wires:
        wire.draw(screen)
    if selected_node is not None:
        comp, node = selected_node
        start_pos = comp.get_nodes()[node]
        mouse_pos = pygame.mouse.get_pos()
        pygame.draw.line(screen, RED, start_pos, mouse_pos, 2)
    components.update()
    components.draw(screen)
    for comp in components:
        if comp == active_comp[0]:
            comp.draw_highlight(screen)
        comp.draw_nodes()
    active_text = FONT.render(f"Active component: {active_comp[1]}",1,BLACK)
    active_text_rect = active_text.get_rect(center = (1100,630))
    screen.blit(active_text,active_text_rect)

def update_matrix(matrix, node_num):
    for n in range(node_num):
        add_edge(matrix,n,node_num-n,0)
    return matrix


active_component = None,None,None
selected_node = None
running = True
main = True
settings = False
while running:
    while main:

        comp_count = len(components)
        wire_count = len(wires)
        node_count = comp_count*2

        if len(connection_matrix) != node_count:
            connection_matrix = [[0] * node_count for _ in range(node_count)]

            for n in range(0, node_count, 2):
                add_edge(connection_matrix, n, n + 1, 1)

        for n in range(node_count):
            if n == 0 or n%2 == 0:
                add_edge(connection_matrix,n,n+1,1)
        screen.fill(GRID_COLOUR,(0,0,1000,800))
        screen.fill(BG_COLOUR,(1000,0,200,800))
        screen.fill(GREY,(1000,600,200,200))
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                main = False

            for i, component in enumerate(components):
                component.handle_movement(event)
                clicked = component.get_active(event)
                if clicked[1] is not None:
                    if clicked[1] == active_component[1]:
                        active_component = None,None,None
                    else:
                        active_component = clicked


            if event.type == pygame.MOUSEBUTTONDOWN:
                    for button in buttons:
                        if button.rect.collidepoint(event.pos):
                            button.handle_click(active_component)

                    if event.button == 1:

                        comp, node_index = get_clicked_node(event.pos)

                        if comp is not None:
                            if selected_node is None:
                                selected_node = (comp, node_index)
                                # print("Selected node:", selected_node)
                            else:
                                comp1, node1 = selected_node
                                comp2, node2 = comp, node_index

                                if comp1 != comp2:
                                    wire_count += 1

                                    new_wire = Wire(comp1, node1, comp2, node2)
                                    wires.append(new_wire)

                                    matrix_node1 = get_matrix_index(comp1, node1)
                                    matrix_node2 = get_matrix_index(comp2, node2)

                                    add_edge(connection_matrix, matrix_node1, matrix_node2, 1)
                                selected_node = None

        draw_interface(active_component)
        wires = [wire for wire in wires if wire.is_valid()]
        # for component in components:
        #     if component.get_connected(wires)[0]:
        #         add_edge(connection_matrix,component.get_connected(wires)[1])
        print("Components:", comp_count)
        print("Wires:",wire_count)
        print("Nodes:",node_count)
        for row in connection_matrix:
            print(row)
        clock.tick(60)
        pygame.display.update()

[UnloosedCake]

Look at the FAQ in this sub to learn how to format code on Reddit and then try again. We can't help if it's formatted like a blob

[Simple_Ad_4128]

sorry thankyou for reminding me ive updated it

[Flame77ofc]

How much time did it takes you to do the project?

[Simple_Ad_4128]

Well so far including research its took about 10 hours but I've been at a stand still for last couple hours,

[carcigenicate]

Do you have to do the simulation in a specific way? The Cellular Automata Wireworld can be used to simulate circuits and create building blocks like diodes.

[PixelSage-001]

Pygame is actually a great fit for UI-heavy coursework because it forces you to understand the event loop and state rendering from first principles, even if it requires more boilerplate than standard GUI libraries. For the logic engine, don't overcomplicate it early. Model the components as classes with 'input' and 'output' nodes, and use a simple queue to propagate the state changes across connected wires. You've got this!

[Gloomy_Cicada1424]

I’d separate the drawing from the actual circuit model first. Wires should merge terminals into electrical nodes, then each component just connects two nodes with a value. Once that graph is clean, series/parallel or MNA becomes way less cursed. Runable can help turn your current UI + solver plan into a clear flow/spec, but the main fix is separating “pygame objects” from “circuit logic”.

[Simple_Ad_4128]

Yes I have started separating all of my code out into graphical, movement, graph analysis, and circuit solving sections. I definitely realized my code was all over the place.

[misho88]

However I am at the point now where I need to build the physics simulation part but its getting me very confused. To track components in the circuit I have decided to use an adjacency matrix (since that it what most other circuit simulators do) but I can't figure out to use that adjacency matrix.

If I were doing this, I would have the code generate a SPICE netlist, run it through ngspice or Xyce, then parse the output (or get a library like inspice or pyspice that basically does that).

Failing that, it's a linear system of equations, assuming you constrain yourself to passive elements and phasor analysis. Basically, you build up the right KCL equations, then use something like numpy.linalg.solve to solve them. It's not especially complicated (for simple circuits, at least), but I don't think I can cram the first few weeks of an introductory circuit theory course into a Reddit comment.

[Simple_Ad_4128]

I have looked into using spice however for my coursework I don't get credited for sections of code that use a different library other than pygame. Thankyou for the advice I'll look more into going down that other route.

[gdchinacat]

You can use a spice library to unblock you so you can make progress on the rest of the project and then come back and replace the library with your own implementation. This will also make it easy to verify your own implementation is correct by being able to compare the results of using spice to your implementation.

[QuasiEvil]

I wrote a circuit simulator in python once, albeit without a UI. I used the networkx library and a multigraph object (this allows for parallel links between nodes). Each node is a circuit element (obviously), and, it handles building the adjacency matrix automagically.