LeGrandChien/legrandchien.py
2026-02-16 15:26:31 +01:00

783 lines
No EOL
25 KiB
Python

from enum import Enum
import numpy as np
from tqdm import tqdm
class Operation(Enum):
VAR = 1
CONST = 2
ADD = 3
SUB = 4
MULT = 5
DIV = 6
POW = 7
SIN = 8
COS = 9
LN = 10
EXP = 11
# const class
class Const:
def __init__(self, value) -> None:
self.value = value
self.type = "Const"
def evaluate(self, dico):
return self.value
def copy(self):
return Const(self.value)
def __add__(self, other):
if other.type == "Const":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) + Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) + Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) + other
def __radd__(self, other):
if isinstance(other, (int, float)):
return Const(other) + self
def __sub__(self, other):
if other.type == "Const":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) - Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) - Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) - other
def __rsub__(self, other):
if isinstance(other, (int, float)):
return Const(other) - self
def __mul__(self, other):
if other.type == "Const":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) * Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) * Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) * other
def __rmul__(self, other):
if isinstance(other, (int, float)):
return Const(other) * self
def __truediv__(self, other):
if other.type == "Const":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) / Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) / Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.CONST,
None,
None,
self.copy()
) / other
def __truediv__(self, other):
if isinstance(other, (int, float)):
return Const(other) / self
# var class
class Var:
def __init__(self, name, degree=0) -> None:
self.name = name
self.type = "Var"
self.degree = degree
def copy(self):
return Var(self.name, self.degree)
def diff(self, n=1):
return Var("d_"*n+self.name, self.degree+n)
def evaluate(self, dico):
return dico[self.name]
def __eq__(self, other):
if isinstance(other, str):
return False
if not (self.type == "Var" and other.type == "Var"):
return False
return self.name == other.name and self.degree == other.degree
def __hash__(self):
return hash((self.name, self.degree))
def __add__(self, other):
if other.type == "Const":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) + Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) + Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) + other
def __radd__(self, other):
if isinstance(other, (int, float)):
return Const(other) + self
def __sub__(self, other):
if other.type == "Const":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) - Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) - Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) - other
def __rsub__(self, other):
if isinstance(other, (int, float)):
return Const(other) - self
def __mul__(self, other):
if other.type == "Const":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) * Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) * Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) * other
def __rmul__(self, other):
if isinstance(other, (int, float)):
return Const(other) * self
def __truediv__(self, other):
if other.type == "Const":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) / Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) / Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(
Operation.VAR,
None,
None,
self.copy()
) / other
def __rtruediv__(self, other):
if isinstance(other, (int, float)):
return Const(other) / self
# equation class
class Equation:
def __init__(self, op, left, right, value):
self.op = op
self.left = left
self.right = right
self.value = value
self.type = "Equation"
self.compiled_function = None
def copy(self):
if self.op == Operation.CONST:
return Equation(Operation.CONST, None, None, Const(self.value.value))
if self.op == Operation.VAR:
return Equation(Operation.VAR, None, None, Var(self.value.name, self.value.degree))
if self.op in {Operation.SIN, Operation.COS}:
return Equation(self.op, None, None, self.value.copy())
return Equation(self.op, self.left.copy(), self.right.copy(), None)
def diff(self):
if self.op == Operation.CONST:
return Equation(Operation.CONST, None, None, Const(0))
if self.op == Operation.VAR:
return Equation(Operation.VAR, None, None, Var("d_"+self.value.name, self.value.degree+1))
if self.op == Operation.ADD:
return self.left.diff() + self.right.diff()
if self.op == Operation.SUB:
return self.left.diff() - self.right.diff()
if self.op == Operation.MULT:
return self.left.diff()*self.right.copy() + self.left.copy()*self.right.diff()
if self.op == Operation.DIV:
return (self.left.diff()*self.right.copy() - self.left.copy()*self.right.diff()) \
/ (self.right.copy()*self.right.copy())
if self.op == Operation.SIN:
return self.value.diff() * Equation(Operation.COS, None, None, self.value.copy())
if self.op == Operation.COS:
return Const(-1) * self.value.diff() * Equation(Operation.SIN, None, None, self.value.copy())
def partial(self, var_name):
if self.op == Operation.CONST:
return Equation(Operation.CONST, None, None, Const(0))
if self.op == Operation.VAR:
if self.value.name == var_name:
return Equation(Operation.CONST, None, None, Const(1))
else:
return Equation(Operation.CONST, None, None, Const(0))
if self.op == Operation.ADD:
return self.left.partial(var_name) + self.right.partial(var_name)
if self.op == Operation.SUB:
return self.left.partial(var_name) - self.right.partial(var_name)
if self.op == Operation.MULT:
return self.left.partial(var_name)*self.right.copy() + self.left.copy()*self.right.partial(var_name)
if self.op == Operation.DIV:
return (self.left.partial(var_name)*self.right.copy() - self.left.copy()*self.right.partial(var_name)) \
/ (self.right.copy()*self.right.copy())
if self.op == Operation.SIN:
return self.value.partial(var_name) * Equation(Operation.COS, None, None, self.value.copy())
if self.op == Operation.COS:
return Const(-1) * self.value.partial(var_name) * Equation(Operation.SIN, None, None, self.value.copy())
def compile(self):
if self.compiled_function is not None:
return self.compiled_function
code = self.generate_code()
self.compiled_function = eval(code)
return self.compiled_function
def generate_code(self):
code, var_count = self._generate_code_recursive(0)
return f"lambda dico: {code}"
def _generate_code_recursive(self, var_count):
if self.op == Operation.CONST:
val = self.value.value
return (str(val), var_count)
elif self.op == Operation.VAR:
val = self.value.name
return (f"dico['{val}']", var_count)
elif self.op == Operation.ADD:
left_code, var_count = self.left._generate_code_recursive(var_count)
right_code, var_count = self.right._generate_code_recursive(var_count)
return (f"({left_code}) + ({right_code})", var_count)
elif self.op == Operation.SUB:
left_code, var_count = self.left._generate_code_recursive(var_count)
right_code, var_count = self.right._generate_code_recursive(var_count)
return (f"({left_code}) - ({right_code})", var_count)
elif self.op == Operation.MULT:
left_code, var_count = self.left._generate_code_recursive(var_count)
right_code, var_count = self.right._generate_code_recursive(var_count)
return (f"({left_code}) * ({right_code})", var_count)
elif self.op == Operation.DIV:
left_code, var_count = self.left._generate_code_recursive(var_count)
right_code, var_count = self.right._generate_code_recursive(var_count)
return (f"({left_code}) / ({right_code})", var_count)
elif self.op == Operation.SIN:
val_code, var_count = self.value._generate_code_recursive(var_count)
return (f"np.sin({val_code})", var_count)
elif self.op == Operation.COS:
val_code, var_count = self.value._generate_code_recursive(var_count)
return (f"np.cos({val_code})", var_count)
return ("0", var_count)
def evaluate(self, dico):
if self.compiled_function is None:
self.compile()
return self.compiled_function(dico)
def solve(self, dico, tmax=10, dt=0.01, progress_bar=True):
self.simplify()
variables = self.getAllVar()
equations = []
# print([v.name for v in variables])
for var in variables:
equations.append(self.partial(var.diff().name).diff() - self.partial(var.name))
unknown = self.getUnknown(dico)
for eq in equations:
eq.simplify()
unknown = unknown.union(eq.getUnknown(dico))
t = 0
values = {t : {var.name : dico[var.name] for var in variables}}
for var in unknown:
dico[var.name] = 0.
# print([var.name for var in unknown])
# print(equations[0])
assert(len(unknown) == len(equations))
n = len(unknown)
variables = list(variables)
unknown = list(unknown)
equations = list(equations)
J_template = np.empty((n, n), dtype=object)
F_template = np.array((n, n), dtype=object)
for i in range(n):
f = equations[i]
f.compile()
F_template[i] = f
for j in range(n):
eq = equations[i].partial(unknown[j].name)
eq.compile()
J_template[i][j] = eq
if progress_bar:
bar = tqdm([n * dt for n in range(int(tmax/dt))])
else:
bar = [n * dt for n in range(int(tmax/dt))]
for t in bar:
err = 1
iterations = 0
max_iteration = 50
while err > 10**-6 and iterations < max_iteration:
iterations += 1
# jacobian
J = np.zeros((n, n))
F = np.zeros(n)
for i in range(n):
F[i] = F_template[i].evaluate(dico)
for j in range(n):
# print(equations[0])
J[i][j] = J_template[i][j].evaluate(dico)
# print(J, F)
err = np.linalg.norm(F)
dX = -np.linalg.inv(J).dot(F)
for i in range(n):
dico[unknown[i].name] += dX[i]
values[t] = {var.name : dico[var.name] for var in variables}
t += dt
# update all values
for i in range(n):
var = unknown[i]
d = var.degree
name = var.name
while d > 0:
ivar_name = name[2:]
dico[ivar_name] += dt * dico[name]
d -= 1
name = ivar_name
return values
def getAllVar(self):
if self.op == Operation.CONST:
return set()
if self.op == Operation.VAR:
if self.value.degree == 0:
return {self.value.copy()}
else:
return {Var(self.value.name[2*self.value.degree:], 0)}
if self.op in {Operation.SIN, Operation.COS}:
return self.value.getAllVar()
return self.left.getAllVar().union(self.right.getAllVar())
def getUnknown(self, dico):
# print(self)
if self.op == Operation.CONST:
return set()
if self.op == Operation.VAR:
if self.value.name not in dico:
return {self.value}
return set()
if self.op in {Operation.SIN, Operation.COS}:
return self.value.getUnknown(dico)
return self.left.getUnknown(dico).union(self.right.getUnknown(dico))
def simplify(self):
if self.left != None:
self.left.simplify()
if self.right != None:
self.right.simplify()
if self.op in {Operation.SIN, Operation.COS}:
self.value.simplify()
if self.op == Operation.CONST:
pass
if self.op == Operation.VAR:
pass
if self.op == Operation.ADD:
if self.left.op == Operation.CONST and self.right.op == Operation.CONST:
self.op = Operation.CONST
self.value = Const(self.left.value.value + self.right.value.value)
self.left = None
self.right = None
elif self.left.op == Operation.CONST and self.left.value.value == 0:
self.op = self.right.op
self.left = self.right.left
self.value = self.right.value
self.right = self.right.right
elif self.right.op == Operation.CONST and self.right.value.value == 0:
self.op = self.left.op
self.right = self.left.right
self.value = self.left.value
self.left = self.left.left
if self.op == Operation.SUB:
if self.left.op == Operation.CONST and self.right.op == Operation.CONST:
self.op = Operation.CONST
self.value = Const(self.left.value.value - self.right.value.value)
self.left = None
self.right = None
elif self.right.op == Operation.CONST and self.right.value.value == 0:
self = self.left
if self.op == Operation.MULT:
if self.left.op == Operation.CONST and self.right.op == Operation.CONST:
self.op = Operation.CONST
self.value = Const(self.left.value.value * self.right.value.value)
self.left = None
self.right = None
elif self.left.op == Operation.CONST and self.left.value.value == 0:
self.op = Operation.CONST
self.left = None
self.right = None
self.value = Const(0)
elif self.right.op == Operation.CONST and self.right.value.value == 0:
self.op = Operation.CONST
self.left = None
self.right = None
self.value = Const(0)
elif self.left.op == Operation.CONST and self.left.value.value == 1:
self.op = self.right.op
self.left = self.right.left
self.value = self.right.value
self.right = self.right.right
elif self.right.op == Operation.CONST and self.right.value.value == 1:
self.op = self.left.op
self.right = self.left.right
self.value = self.left.value
self.left = self.left.left
if self.op == Operation.SIN:
if self.value.op == Operation.CONST:
self.op = Operation.CONST
self.left = None
self.right = None
self.value = Const(np.sin(self.value.value))
if self.op == Operation.COS:
if self.value.op == Operation.CONST:
self.op = Operation.CONST
self.left = None
self.right = None
self.value = Const(np.cos(self.value.value))
def __str__(self):
if self.op == Operation.CONST:
return str(self.value.value)
if self.op == Operation.VAR:
return self.value.name
if self.op == Operation.ADD:
return f"{str(self.left)} + {str(self.right)}"
if self.op == Operation.SUB:
return f"{str(self.left)} - {str(self.right)}"
if self.op == Operation.MULT:
return f"({str(self.left)})*({str(self.right)})"
if self.op == Operation.DIV:
return f"({str(self.left)})/({str(self.right)})"
if self.op == Operation.SIN:
return f"sin({str(self.value)})"
if self.op == Operation.COS:
return f"cos({str(self.value)})"
def __add__(self, other):
if other.type == "Const":
return self + Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return self + Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(Operation.ADD, self, other, None)
def __radd__(self, other):
if isinstance(other, (int, float)):
return Const(other) + self
def __sub__(self, other):
if other.type == "Const":
return self - Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return self - Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(Operation.SUB, self, other, None)
def __rsub__(self, other):
if isinstance(other, (int, float)):
return Const(other) - self
def __mul__(self, other):
if other.type == "Const":
return self * Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return self * Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(Operation.MULT, self, other, None)
def __rmul__(self, other):
if isinstance(other, (int, float)):
return Const(other) * self
def __truediv__(self, other):
if other.type == "Const":
return self / Equation(
Operation.CONST,
None,
None,
other
)
if other.type == "Var":
return self / Equation(
Operation.VAR,
None,
None,
other
)
if other.type == "Equation":
return Equation(Operation.DIV, self, other, None)
def __rtruediv__(self, other):
if isinstance(other, (int, float)):
return Const(other) / self
def cos(value):
if isinstance(value, (float, int)):
return Equation(Operation.CONST, None, None, np.cos(value))
if value.type == "Const":
return Equation(Operation.CONST, None, None, np.cos(value.value))
if value.type == "Var":
return Equation(Operation.COS, None, None, Equation(Operation.VAR, None, None, value.copy()))
if value.type == "Equation":
return Equation(Operation.COS, None, None, value.copy())
def sin(value):
if isinstance(value, (float, int)):
return Equation(Operation.CONST, None, None, np.sin(value))
if value.type == "Const":
return Equation(Operation.CONST, None, None, np.sin(value.value))
if value.type == "Var":
return Equation(Operation.SIN, None, None, Equation(Operation.VAR, None, None, value.copy()))
if value.type == "Equation":
return Equation(Operation.SIN, None, None, value.copy())