Merge pull request #9 from alexanderthclark/eqm

Add equilibrium class

microecon/curves.py

@@ -443,7 +443,7 @@ def midpoint_elasticity(self, p1, p2):
         return self.price_elasticity(mean_p)
 
 
-    def plot(self, ax = None, textbook_style = True, max_q = 10,
+    def plot(self, ax = None, textbook_style = True, max_q = None,
              color = 'black', linewidth = 2, label = True):
         """
         Plot the supply or demand curve.
@@ -476,13 +476,15 @@ def plot(self, ax = None, textbook_style = True, max_q = 10,
         # core plot
         if self._domain:
             x1, x2 = self._domain
+            if x2 == np.inf:
+                x2 = max_q if max_q else x1*2
         else:
             x1 = 0
             q_ = self.q_intercept
             if np.isnan(q_): # if slope is 0
                 q_ = 10 ** 10
             if type(self).__name__ == "Supply":
-                x2 = np.max([max_q, q_*2])
+                x2 = np.max([10, q_*2])
             else:
                 x2 = q_
 
@@ -669,7 +671,13 @@ def __init__(self, intercept=None, slope=None, elements=None, inverse = True):
         sections = [(cuts[i], cuts[i+1]) for i in range(len(cuts)-1)]
         qsections = [ (self.q(ps[0]), self.q(ps[1]))  for ps in sections]
         sections.append( (cuts[-1], np.inf) )
-        qsections.append((0,0))
+        if self.q(cuts[-1]+1) <= 0: # demand
+            qsections.append((0,0))
+        elif len(qsections): # supply
+            maxq = np.max(qsections[-1])
+            qsections.append((maxq, np.inf))
+        else: # supply
+            qsections.append((0, np.inf))
         self.psections = sections
         self.qsections = qsections
         self._set_piece_domains()
@@ -735,7 +743,7 @@ def from_formula(cls, equation: str):
 
     def __call__(self, x):
         """
-        Computes p given q=x. This is wrong currently.
+        Computes p given q=x.
 
         Parameters
         ----------
@@ -745,7 +753,15 @@ def __call__(self, x):
         -------
         float
         """
-        return np.sum([np.max([0, c(x)]) for c in self.elements])
+
+        for piece in self.pieces:
+            if piece:
+                a, b = piece._domain
+                if (a <= x <= b) or (a >= x >= b):
+                    return piece(x)
+        # might be x out of limits
+        return np.nan
+        #return np.sum([np.max([0, c(x)]) for c in self.elements])
 
     def q(self, p):
         # returns q given p
@@ -774,21 +790,43 @@ def __add__(self, other):
         elements = self.elements + other.elements
         return Affine(elements=elements)
 
-    def plot(self, ax=None):
+    def plot(self, ax=None, set_lims=True, max_q=None):
+        '''
+        wip
 
+        follow this up with an equilibrium plot
+        '''
         if ax is None:
             fig, ax = plt.subplots()
 
         # use intersections to 
         for piece in self.pieces:
             # get two points
             if piece:
-                piece.plot(ax=ax, label=False)
-
-        ax.set_ylim(0, np.max(self.intercept)*1.01)
+                piece.plot(ax=ax, label=False, max_q=max_q)
+
+        # check limits
+        if set_lims:
+            ylim = ax.get_ylim()
+            if ylim[1] <= np.max(self.intercept):
+                ax.set_ylim(0, np.max(self.intercept)*1.01)
+
+            flat_q = sorted([i for tup in self.qsections for i in tup])
+            flat_p = sorted([i for tup in self.psections for i in tup])
+            if max_q is None:
+                if np.inf in flat_q:
+                    max_q = 1.5*flat_q[-2]
+                else:
+                    max_q = flat_q[-1]
+            if np.inf in flat_p:
+                max_p = np.max([self(max_q), 1.5*flat_p[-2]])
+            else:
+                max_p = np.max([self(max_q), 1.5*flat_p[-1]])
 
-        # fix for demand and supply and inverse vs q(p)
-        #ax.set_xlim(0, np.max(self.intercept))
+            ax.set_ylim(0, max_p)
+            ax.set_xlim(0, max_q)
+            # fix for demand and supply and inverse vs q(p)
+            #ax.set_xlim(0, np.max(self.intercept))
 
 
 class Demand(Affine):
@@ -797,7 +835,6 @@ def __init__(self, intercept=None, slope=None, elements=None, inverse = True):
         """
         Initializes a Demand curve object.
         """
-
         super().__init__(intercept, slope, elements, inverse)
         self._check_slope()
 
@@ -812,7 +849,6 @@ def __init__(self, intercept=None, slope=None, elements=None, inverse = True):
         """
         Initializes a Supply curve object.
         """
-
         super().__init__(intercept, slope, elements, inverse)
         self._check_slope()
 

microecon/equilibrium.py

@@ -0,0 +1,31 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from microecon.curves import Demand, Supply, intersection
+
+class Equilibrium:
+
+    def __init__ (self, demand: Demand, supply: Supply):
+        """Equilibrium produced by demand and supply. Initialized as market equilibrium."""
+
+        self.demand = demand
+        self.supply = supply
+        for dpiece in demand.pieces:
+            for spiece in supply.pieces:
+                if dpiece and spiece:
+                    s_domain = spiece._domain
+                    d_domain = dpiece._domain
+                    p, q = intersection(dpiece, spiece)
+                    in_demand = d_domain[1] <= q <= d_domain[0]
+                    in_supply = s_domain[0] <= q <= s_domain[1]
+                    if in_demand and in_supply:
+                        self.q = q
+                        self.p = p
+                        break
+
+    def plot(self):
+
+        q_int = np.max([piece.q_intercept for piece in self.demand.pieces if piece])
+
+        fig, ax = plt.subplots()
+        self.demand.plot(ax=ax)
+        self.supply.plot(ax=ax, max_q=q_int)