def AddVariable(self, varname, desc='', eqn=''):
"""
Add a variable to the sector.
The variable name (varname) is the local name; it will be decorated to create a
full name. Equations within a sector can use the local name; other sectors need to
use GetVariableName to get the full name.
:param varname: str
:param desc: str
:param eqn: str
:return: None
"""
if '__' in varname:
raise ValueError('Cannot use "__" inside local variable names: ' +
varname)
if desc is None:
desc = ''
if type(eqn) == Equation:
equation = eqn
else:
equation = Equation(varname, desc, [
Term(eqn, is_blob=True),
])
if varname in self.GetVariables():
Logger('[ID={0}] Variable Overwritten: {1}',
priority=3,
data_to_format=(self.ID, varname))
self.EquationBlock.AddEquation(equation)
# self.Equations[varname] = eqn
Logger('[ID={0}] Variable Added: {1} = {2} # {3}',
priority=2,
data_to_format=(self.ID, varname, eqn, desc))
def test_list(self):
block = EquationBlock()
eq = Equation('x', rhs=[Term('y')])
block.AddEquation(eq)
block.AddEquation(Equation('a'))
# Always sorted
self.assertEqual(['a', 'x'], block.GetEquationList())
def test_str(self):
t = Term('x')
t.Constant = 0.
self.assertEqual('', str(t))
t.Constant = 1.
self.assertEqual('+x', str(t))
t.Constant = -1.
self.assertEqual('-x', str(t))
t.Constant = 2.
self.assertEqual('+{0}*x'.format(str(2.)), str(t))
t.Constant = -2.
self.assertEqual('-{0}*x'.format(str(2.)), str(t))
def AddCashFlow(self, term, eqn=None, desc=None, is_income=True):
"""
Add a cash flow to the sector. Will add to the financial asset equation (F), and
the income equation (INC) if is_income is True.
Except: There is a list of exclusions to which cash flows are not considered income.
That setting will override the is_income parameter. This allows us to carve out exceptions
to the standard behaviour, which generally is to assume that cash flows are associated with
income.
:param term: str
:param eqn: str
:param desc: str
:param is_income: bool
:return: None
"""
term = term.strip()
if len(term) == 0:
return
term_obj = Term(term)
if not term_obj.IsSimple: # pragma: no cover - Not implemented; cannot hit the line below.
raise LogicError(
'Must supply a single variable as the term to AddCashFlow')
# term = term.replace(' ', '')
# if not (term[0] in ('+', '-')):
# term = '+' + term
# if len(term) < 2:
# raise ValueError('Invalid cash flow term')
self.EquationBlock['F'].AddTerm(term)
if is_income:
# Need to see whether it is excluded
mod = self.GetModel()
for obj, excluded in mod.IncomeExclusions:
if obj.ID == self.ID:
if term_obj.Term == excluded:
is_income = False
break
if is_income:
self.EquationBlock['INC'].AddTerm(term)
if eqn is None:
return
# Remove the +/- from the term
term = term_obj.Term
if term in self.GetVariables():
rhs = self.EquationBlock[term].RHS()
if rhs == '' or rhs == '0.0':
self.SetEquationRightHandSide(term, eqn)
else:
self.AddVariable(term, desc, eqn)
def AddTermToEquation(self, varname, term):
"""
Add a new term to an existing equation.
The term variable may be either a string or (non-Blob) Term object.
:param varname: str
:param term: Term
:return: None
"""
term = Term(term)
Logger('Adding term {0} to Equation {1} in Sector {2} [ID={3}]',
priority=2,
data_to_format=(term, varname, self.Code, self.ID))
try:
self.EquationBlock[varname].AddTerm(term)
except KeyError:
raise KeyError('Variable {0} not in Sector {1}'.format(
varname, self.Code))
def SetEquationRightHandSide(self, varname, rhs):
"""
Set the right hand side of the equation for an existing variable.
:param varname: str
:param rhs: str
:return: None
"""
try:
self.EquationBlock[varname].TermList = [
Term(rhs, is_blob=True),
]
except KeyError:
raise KeyError('Variable {0} does not exist'.format(varname))
# Could try: Equation.ParseString(rhs), but is too slow in unit tests...
# if varname not in self.Equations:
# raise KeyError('Variable {0} does not exist'.format(varname))
Logger('[ID={0}] Equation set: {1} = {2} ',
priority=2,
data_to_format=(self.ID, varname, rhs))
def test_copy(self):
t = Term('x')
t2 = Term(t)
self.assertEqual('x', t2.Term)
def test_simple_3(self):
t = Term('x')
self.assertTrue(t.IsSimple)
self.assertEqual('x', t.Term)
def test_simple_2(self):
# Original comment: constants not supported yet as "Simple"
# 2019-05-04: Not sure why no support for constants?
t = Term('2')
self.assertTrue(t.IsSimple)
self.assertEqual('2', t.Term)
def test_interior_2(self):
with self.assertRaises(LogicError):
t = Term('-(y+x)')
def test_sign_4(self):
t = Term('-(+x)')
self.assertEqual(-1.0, t.Constant)
def test_sign_3(self):
t = Term('-(-x)')
self.assertEqual(1.0, t.Constant)
def test_sign_2(self):
t = Term(' -x')
self.assertEqual(-1.0, t.Constant)
def test_multiply(self):
t = Term('a*b')
self.assertEqual('+a*b', str(t))
def test_access(self):
block = EquationBlock()
eq = Equation('x', rhs=[Term('y')])
block.AddEquation(eq)
out = block['x']
self.assertEqual('y', out.RHS())
def test_AddTermFail(self):
eq = Equation('x', 'desc', 'y')
t2 = Term('y^2', is_blob=True)
with self.assertRaises(LogicError):
eq.AddTerm(t2)
def test_str_2(self):
t1 = Term('y')
t2 = Term('-z')
eq = Equation('x', 'define x', (t2, t1))
self.assertEqual('-z+y', eq.GetRightHandSide())
def bad_ctor(self):
t1 = Term('x')
t2 = Term('y^2', is_blob=True)
with self.assertRaises(LogicError):
Equation('lhs', '', (t1, t2))
def test_str_1(self):
t1 = Term('y')
t2 = Term('-z')
eq = Equation('x', 'define x', (t1, t2))
self.assertEqual('y-z', eq.GetRightHandSide())
def test_simple_2(self):
# NOTE: constants not supported yet as "Simple"
with self.assertRaises(NotImplementedError):
t = Term('2')
def test_blob(self):
t = Term('(bazoonga*kablooie)^2', is_blob=True)
self.assertEqual('(bazoonga*kablooie)^2', str(t))
def test_fail_bracket(self):
with self.assertRaises(SyntaxError):
Term('+(x ')
def test_str_not_simple(self):
with self.assertRaises(NotImplementedError):
t = Term('f(x)')
def test_sign_1(self):
t = Term('x')
self.assertEqual(1.0, t.Constant)
