<s> import numpy as np <EOL> def view ( A , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> from numpy . lib . stride_tricks import as_strided <EOL> assert A . ndim == <NUM_LIT:2> , "<STR_LIT>" <EOL> assert A . shape [ <NUM_LIT:0> ] == A . shape [ <NUM_LIT:1> ] , "<STR_LIT>" <EOL> if offset > <NUM_LIT:0> : <EOL> return as_strided ( A [ <NUM_LIT:0> , offset : ] , shape = ( A . shape [ <NUM_LIT:0> ] - offset , ) , strides = ( ( A . shape [ <NUM_LIT:0> ] + <NUM_LIT:1> ) * A . itemsize , ) ) <EOL> elif offset < <NUM_LIT:0> : <EOL> return as_strided ( A [ - offset : , <NUM_LIT:0> ] , shape = ( A . shape [ <NUM_LIT:0> ] + offset , ) , strides = ( ( A . shape [ <NUM_LIT:0> ] + <NUM_LIT:1> ) * A . itemsize , ) ) <EOL> else : <EOL> return as_strided ( A , shape = ( A . shape [ <NUM_LIT:0> ] , ) , strides = ( ( A . shape [ <NUM_LIT:0> ] + <NUM_LIT:1> ) * A . itemsize , ) ) <EOL> def offdiag_view ( A , offset = <NUM_LIT:0> ) : <EOL> from numpy . lib . stride_tricks import as_strided <EOL> assert A . ndim == <NUM_LIT:2> , "<STR_LIT>" <EOL> Af = as_strided ( A , shape = ( A . size , ) , strides = ( A . itemsize , ) ) <EOL> return as_strided ( Af [ ( <NUM_LIT:1> + offset ) : ] , shape = ( A . shape [ <NUM_LIT:0> ] - <NUM_LIT:1> , A . shape [ <NUM_LIT:1> ] ) , strides = ( A . strides [ <NUM_LIT:0> ] + A . itemsize , A . strides [ <NUM_LIT:1> ] ) ) <EOL> def _diag_ufunc ( A , b , offset , func ) : <EOL> b = np . squeeze ( b ) <EOL> assert b . ndim <= <NUM_LIT:1> , "<STR_LIT>" <EOL> dA = view ( A , offset ) ; func ( dA , b , dA ) <EOL> return A <EOL> def times ( A , b , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return _diag_ufunc ( A , b , offset , np . multiply ) <EOL> multiply = times <EOL> def divide ( A , b , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return _diag_ufunc ( A , b , offset , np . divide ) <EOL> def add ( A , b , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return _diag_ufunc ( A , b , offset , np . add ) <EOL> def subtract ( A , b , offset = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return _diag_ufunc ( A , b , offset , np . subtract ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> import doctest <EOL> doctest . testmod ( ) </s>
