Add sma indicator

data/indicators/sma.py

@@ -0,0 +1,21 @@
+def calculate_sma(data, period):
+    """
+    Calculate the Simple Moving Average (SMA) for the given data.
+    :param data: A list of dictionaries containing 'date' and 'close' prices.
+    :param period: period used to calculate sma.
+    :return: A list of dictionaries with an additional 'SMA' key.
+    """
+    close_prices = [entry['close'] for entry in data]
+    sma_values = []
+
+    for i in range(len(close_prices)):
+        if i + 1 < period:
+            sma_values.append(None)
+        else:
+            sma = sum(close_prices[i + 1 - period:i + 1]) / period
+            sma_values.append(sma)
+
+    for i in range(len(data)):
+        data[i]['SMA'] = sma_values[i]
+
+    return data

database/__init__.py

@@ -2,6 +2,8 @@ from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.orm import scoped_session
 from model.base import Base
+from sqlalchemy import MetaData, Table
+
 
 
 class Database:
@@ -16,3 +18,11 @@ class Database:
 
     def drop_tables(self):
         Base.metadata.drop_all(self.engine)
+
+
+    def drop_table(self, table):
+        meta = MetaData()
+        meta.reflect(bind=self.engine)
+        if table in meta.tables:
+            table = Table(table, meta, autoload_with=self.engine)
+            table.drop(self.engine, checkfirst=True)

manage.py

@@ -1,6 +1,9 @@
 import csv
+import copy
 from datetime import datetime
 from database import Database
+from data.indicators.sma import calculate_sma
+from model.sma import Sma
 from model.ticker import Ticker
 
 
@@ -46,5 +49,45 @@ def load_database(chunk_size=100000):
             pprint.pprint(e)
 
 
+def create_sma():
+    from sqlalchemy import distinct
+
+    database.drop_table("sma")
+    database.create_tables()
+
+    sma_periods = [5, 10, 20, 50, 100, 200]
+    symbols = database.session.query(distinct(Ticker.symbol)).all()
+    count = 0
+    for symbol in symbols:
+        try:
+            daily_ticker_data = database.session.query(Ticker).filter_by(symbol=symbol[0]).all()
+            daily_ticker_data = [ticker.to_dict() for ticker in daily_ticker_data]
+            
+            smas = {}
+            for period in sma_periods:
+                data = copy.deepcopy(daily_ticker_data)
+                sma_data = calculate_sma(data, period)
+                smas[period] = sma_data
+            
+            for i in range(len(daily_ticker_data)):
+                row = Sma(
+                    date=daily_ticker_data[i]['date'],
+                    symbol=symbol[0],
+                    sma_5=smas[5][i]['SMA'],
+                    sma_10=smas[10][i]['SMA'],
+                    sma_20=smas[20][i]['SMA'],
+                    sma_50=smas[50][i]['SMA'],
+                    sma_100=smas[100][i]['SMA'],
+                    sma_200=smas[200][i]['SMA']
+                )
+                database.session.add(row)
+            database.session.commit()
+            count += 1
+            print(f"finished {symbol[0]}: {count} of {len(symbols)}")
+        except Exception as e:
+            print(e)
+            database.session.rollback()
+
+
 if __name__ == "__main__":
-    load_database()
+    create_sma()

model/sma.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+from sqlalchemy import Column
+from sqlalchemy import DateTime
+from sqlalchemy import Float
+from sqlalchemy import Integer
+from sqlalchemy import String
+from model.base import Base
+
+
+class Sma(Base):
+    __tablename__ = 'sma'
+
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    date = Column(DateTime)
+    symbol = Column(String(15), nullable=False)
+    sma_5 = Column(Float)
+    sma_10 = Column(Float)
+    sma_20 = Column(Float)
+    sma_50 = Column(Float)
+    sma_100 = Column(Float)
+    sma_200 = Column(Float)

model/ticker.py

@@ -19,3 +19,15 @@ class Ticker(Base):
     low = Column(Float)
     close = Column(Float)
     volume = Column(Integer)
+
+
+    def to_dict(self):
+        return {
+            "date": self.date,
+            "symbol": self.symbol,
+            "open": self.open,
+            "high": self.high,
+            "low": self.low,
+            "close": self.close,
+            "volume": self.volume,
+        }