NameError: name 'cos' is not defined, sageMath9.4

I am writing a python module using SageMath9.4. Basically, I want to import this module into juypterlab notebooks (running a SageMath 9.4 kernel) to do calculations etc.

Here is the start of it:

class Coxeter_System:
'''This class defines the standard root system associated to an abstract Coxeter group.'''

def __init__(self, coxeter_matrix):
    '''Sets up a Coxeter system and root system. At this stage, limited to up to rank 7. 
    '''
    def set_up_coefficient_space(coxeter_matrix):
        '''Sets up a polynomial ring a Free module over a polynomial ring quotiented by the minimal polynomials of the 
        non-rational cos(pi/m_ij) values. 
        
        This is so roots can be compared using an abstract free module rather than over reals'''
        
        A = coxeter_matrix
        k = len(A.rows())

        # Get the cos(pi/m_ij) which are irrational
        non_rational_angles = [x for x in [cos(pi/x) for x in set(A[i,j] for i in range(0,k) for j in range(0,k))] if x not in QQ]

However, when I open another Juypterlab session, import the python module and try to create an instance of the object "Coxeter_System", I get the following error (I have tried to do from math import cos both from the notebook where I want to import the module to, and in the module itself, but I still get the same error.

Any help would be greatly appreciated!):

NameError                                 Traceback (most recent call last)
<ipython-input-6-6b542b6cb042> in <module>
----> 1 W = c.Coxeter_System(Matrix([[Integer(1),Integer(3),Integer(4)],[Integer(3),Integer(1),Integer(3)],[Integer(4),Integer(3),Integer(1)]]))

~/coxeter_groups.py in __init__(self, coxeter_matrix)
     60             return matrix(R,k,B)
     61 
---> 62         R = set_up_coefficient_space(coxeter_matrix)
     63         A = coxeter_matrix
     64         k = len(A.rows())

~/coxeter_groups.py in set_up_coefficient_space(coxeter_matrix)
     17 
     18             # Get the cos(pi/m_ij) which are irrational
---> 19             non_rational_angles = [x for x in [cos(pi/x) for x in set(A[i,j] for i in range(0,k) for j in range(0,k))] if x not in QQ]
     20 
     21             # sort the irrational values of cos(pi/m_ij) in ascending order

~/coxeter_groups.py in <listcomp>(.0)
     17 
     18             # Get the cos(pi/m_ij) which are irrational
---> 19             non_rational_angles = [x for x in [cos(pi/x) for x in set(A[i,j] for i in range(0,k) for j in range(0,k))] if x not in QQ]
     20 
     21             # sort the irrational values of cos(pi/m_ij) in ascending order

NameError: name 'cos' is not defined

from NameError: name 'cos' is not defined, sageMath9.4