Beginnings of a neuron simulation

Question

I have 2 classes that do extremely similar things. The only differences are that they draw using different colors, and that one draws in the back.

The classes NeuronOutput and NeuronReciever are very similar.

class NeuronOutput():
 """The place where a Neuron outputs its signal"""
 def __init__(self,neuron):
 self.neuron=neuron
 width=self.neuron.width*0.5
 offset=width/2
 distance=self.neuron.width*0.4
 self.pos=[self.neuron.x+math.sin(self.neuron.angle)*distance,self.neuron.y-math.cos(self.neuron.angle)*distance]
 self.drawRect=[self.pos[0]-offset,self.pos[1]-offset,width,width] 
 def draw(self,surf):
 pygame.draw.ellipse(surf,[200,150,200],self.drawRect)
class NeuronReciever():
 """Where a Neuron recieves input"""
 def __init__(self,neuron):
 self.neuron=neuron
 width=self.neuron.width*0.5
 offset=width/2
 distance=self.neuron.width*0.4
 self.pos=[self.neuron.x-math.sin(self.neuron.angle)*distance,self.neuron.y+math.cos(self.neuron.angle)*distance]
 self.drawRect=[self.pos[0]-offset,self.pos[1]-offset,width,width] 
 def draw(self,surf):
 pygame.draw.ellipse(surf,[200,255,200],self.drawRect)

There is also a class Neuron that uses instances of the two above classes as properties. It is also fairly similar to them, however it is a bit different in the way it initializes and draws.

class Neuron():
 """A single neuron with Synapses that connect to others"""
 def __init__(self,x,y):
 self.x=x
 self.y=y
 self.pos=[x,y]
 self.width=30
 self.drawOffset=self.width/2;
 self.connections=[]
 self.angle=math.radians(random.randint(0,360))
 self.drawRect=[self.x-self.drawOffset,self.y-self.drawOffset,self.width,self.width]
 self.output=NeuronOutput(self)
 self.input=NeuronReciever(self)
 def draw(self,surf):
 pygame.draw.ellipse(surf,[200,150,50],self.drawRect)
 self.output.draw(surf)
 self.input.draw(surf)
 self.syn.draw(surf)

(The code is the beginning of a neuron simulation.)

I want to keep all these as separate classes for the sake of readability and ease of use. How can I reduce the repetitiveness of this code, while still keeping it readable?

Answer 1

So you have two classes which are very similar and differ only in one thing? Then make a parent class that has all the behavior and let the others inherit from it or add a switch to flip between the two behaviors.

I would probably go with the first. To accomplish this, make self.drawRect a property (and rename it to self.rect, it does not actually draw anything, it just says where to draw):

class NeuronIO:
 """Draw input or output of a neuron.
 Abstract base class.
 """
 color = 0, 0, 0 # default color is black
 def __init__(self, neuron):
 self.neuron = neuron
 width = 0.5 * self.neuron.width
 offset = width / 2
 distance = 0.4 * self.neuron.width
 self._pos = None
 self._rect = None
 def draw(self, surf):
 pygame.draw.ellipse(surf, self.color, self.rect)
 @property
 def pos(self):
 """Needs to be implemented by the child classes"""
 raise NotImplementedError
 @property
 def rect(self):
 if self._rect is None:
 pos = self.pos
 self._rect = pos[0] - offset, pos[1] - offset, width, width 
 return self._rect
class NeuronOutput(NeuronIO):
 """The place where a Neuron outputs its signal"""
 color = 200, 150, 200
 @property
 def pos(self):
 if self._pos is None:
 self._pos = (self.neuron.x + math.sin(self.neuron.angle) * distance,
 self.neuron.y - math.cos(self.neuron.angle) * distance)
 return self._pos
class NeuronReciever(NeuronIO):
 """Where a Neuron receives input"""
 color = 200, 255, 200
 @property
 def pos(self):
 if self._pos is None:
 self._pos = (self.neuron.x - math.sin(self.neuron.angle) * distance,
 self.neuron.y + math.cos(self.neuron.angle) * distance)
 return self._pos

Here both rect and pos are cached so that they are only calculated the first time they are needed.

In addition I added some whitespace around operators and after commas, as recommended by Python's official style-guide, PEP8.

---

At this point you should ask yourself: "But how often do I need to draw the receiver or output without also drawing a neuron?". If the answer to that question is "probably never!", then you should just fold all of that functionality into the Neuron class:

import math
import random
import pygame
class Neuron:
 """A single neuron with Synapses that connect to others"""
 def __init__(self, x, y):
 self.x = x
 self.y = y
 self.pos = x, y
 self.width = 30
 self.offset = self.width / 2
 self.connections = []
 self.angle = math.radians(random.randint(0, 360))
 self.rect = (self.x - self.offset, self.y - self.offset,
 self.width, self.width)
 self.color = 200, 150, 50
 def draw(self, surf):
 pygame.draw.ellipse(surf, self.color, self.rect)
 pygame.draw.ellipse(surf, *self.input)
 pygame.draw.ellipse(surf, *self.output)
 @property
 def input(self):
 color = 200, 255, 200
 rect = self.calc_rect()
 return color, rect
 @property
 def output(self):
 color = 200, 255, 200
 rect = self.calc_rect(output=True)
 return color, rect
 def calc_rect(self, output=False):
 width = 0.5 * self.width
 distance = 0.4 * self.width
 offset = 0.25 * self.width
 if output:
 pos = (self.x + math.sin(self.angle) * distance,
 self.y - math.cos(self.angle) * distance)
 else:
 pos = (self.x - math.sin(self.angle) * distance,
 self.y + math.cos(self.angle) * distance)
 return pos[0] - offset, pos[1] - offset, width, width

Here I did not go through the effort of making sure that Neuron.input and Neuron.output are only calculated once, but you could do it in the same way as I did above.

Note that I made all of your lists into tuples (like Neuron.pos). This is just aesthetics atm, but if you want to be able to do e.g. neuron.pos[0] += 1, then you need to make them lists again.

• Luke B: Thanks, that makes sense and is quite helpful! p.s. you missed that they draw in different colors. color is the second argument of pygame.draw.ellipse
• Graipher: @pydude Fixed, made the color custom as well.
• Graipher: @pydude And I added some code for the case that this should not actually be standalone classes.