python-utils/matrix.py at main · theodor8/python-utils · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146


class Mat:
    def __init__(self, m):
        self.m = m

    def size(self):
        return len(self.m), len(self.m[0]) # rows, cols

    def transpose(self):
        rs, cs = self.size()
        return Mat([[self.m[r][c]  for r in range(rs)] for c in range(cs)])

    def col(self, i):
        rs, _ = self.size()
        return [self.m[r][i] for r in range(rs)]

    def copy(self):
        return Mat([[x for x in r] for r in self.m])

    def flip(self):
        return Mat([r[::-1] for r in self.m[::-1]])

    @staticmethod
    def unit(n):
        return Mat([[int(i == j) for j in range(n)] for i in range(n)])

    @staticmethod
    def all(v, n, m):
        return Mat([[v for _ in range(m)] for _ in range(n)])

    @staticmethod
    def zeros(n, m):
        return Mat.all(0, n, m)

    @staticmethod
    def ones(n, m):
        return Mat.all(1, n, m)

    def __repr__(self):
        s = ''
        for r in self.m:
            s += f"[{' '.join(str(x) for x in r)}]\n"
        return s[:-1]

    def __add__(self, other):
        ar, ac = self.size()
        br, bc = other.size()
        if ar != br or ac != bc:
            raise ValueError('size(a) != size(b)')
        return Mat([[self.m[r][c] + other.m[r][c] for c in range(ac)] for r in range(ar)])

    def __sub__(self, other):
        ar, ac = self.size()
        br, bc = other.size()
        if ar != br or ac != bc:
            raise ValueError('size(a) != size(b)')
        return Mat([[self.m[r][c] - other.m[r][c] for c in range(ac)] for r in range(ar)])

    def __mul__(self, other):
        if isinstance(other, (int, float)):
            return Mat([[x * other for x in r] for r in self.m])
        ar, ac = self.size()
        br, bc = other.size()
        if ac != br:
            raise ValueError('cols(a) != rows(b)')
        return Mat([[dot(self.m[r], other.col(c)) for c in range(bc)] for r in range(ar)])

    def __rmul__(self, other):
        return self * other

    def __pow__(self, other):
        if not isinstance(other, int):
            raise ValueError('exponent must be int')
        if other < 0:
            raise ValueError('negative exponent')
        rs, cs = self.size()
        if rs != cs:
            raise ValueError('Rows != Cols')
        if other == 0:
            return Mat.unit(rs)
        m = self.copy()
        for _ in range(other - 1):
            m *= self
        return m

    def invert(self): # gauss
        m = self.copy()
        rs, cs = m.size()
        if rs != cs:
            raise ValueError('Rows != Cols')
        n = rs
        b = Mat.unit(n)
        m, b = gauss(m, b)
        m, b = gauss(m.flip(), b.flip())
        return b.flip()

class ColVec(Mat):
    def __init__(self, *v: float):
        super().__init__([[x] for x in v])

    def __repr__(self):
        s = ''
        for r in self.m:
            s += f"[{r[0]}]\n"
        return s[:-1]

class RowVec(Mat):
    def __init__(self, *v: float):
        super().__init__([v])

    def __repr__(self):
        s = ''
        for r in self.m:
            s += f"[{' '.join(str(x) for x in r)}]\n"
        return s[:-1]


def dot(a, b):
    return sum(i * j for i, j in zip(a, b))


def gauss(m, b):
    m = m.copy()
    b = b.copy()
    mr, mc = m.size()
    br, bc = b.size()
    if mr != br:
        raise ValueError('mr != br')
    n = min(mr, mc)
    for i in range(n):
        for j in range(i + 1, n):
            s = -m.m[j][i] / m.m[i][i] # scalar
            for k in range(mc):
                m.m[j][k] += m.m[i][k] * s
            for k in range(bc):
                b.m[j][k] += b.m[i][k] * s
    for i in range(n):
        s = m.m[i][i]
        for k in range(mc):
            m.m[i][k] /= s
        for k in range(bc):
            b.m[i][k] /= s
    return m, b