--- a/pydct/dct.py	Tue May 06 23:01:28 2008 +0200
+++ b/pydct/dct.py	Tue May 06 23:39:54 2008 +0200
@@ -26,19 +26,35 @@
     c = __init_dctN(len(v))
     return numpy.dot(numpy.transpose(c), v)
 
-def zigzag(x):
+import dct, numpy
+
+def fdctNxN(b):
+    '''
+    Forward 2D DCT on NxN array.
+    '''
+    c = __init_dctN(b.shape[0])
+    return numpy.dot(c, numpy.dot(b, numpy.transpose(c)))
+
+def idctNxN(b):
+    '''
+    Inverse 2D DCT on NxN array.
+    '''
+    c = __init_dctN(b.shape[0])
+    return numpy.dot(numpy.transpose(c), numpy.dot(b, c))
+
+def zigzag(x, start=0, stop=None):
     """Generates zig-zag scan sequence for two dimensional array."""
-    j = 0
-    i = 0
-    while True:
-
+    i = j = count = 0
+    if not stop: stop = x.size
+    while count < stop:
         if i == x.shape[0]: j += 2; i = x.shape[0]-1
         else: j = 0
         while i >= 0 and j < x.shape[1]:
             # run up
-            yield i, j
+            if start <= count < stop: yield i, j
             i -= 1
             j += 1
+            count += 1
 
         if j > x.shape[1]: break
 
@@ -46,9 +62,10 @@
         else: i = 0
         while j >= 0 and i < x.shape[0]:
             # run down
-            yield i, j
+            if start <= count < stop: yield i, j
             i += 1
             j -= 1
+            count += 1
 
         if i > x.shape[0]: break
 

--- a/pydct/tNxN.py	Tue May 06 23:01:28 2008 +0200
+++ b/pydct/tNxN.py	Tue May 06 23:39:54 2008 +0200
@@ -4,12 +4,13 @@
     '''
     Forward 2D DCT on NxN array.
     '''
-    c = __init_dctN(b.shape[0])
+    c = dct.__init_dctN(b.shape[0])
     return numpy.dot(c, numpy.dot(b, numpy.transpose(c)))
 
 def idctNxN(b):
     '''
     Inverse 2D DCT on NxN array.
     '''
-    c = __init_dctN(b.shape[0])
+    c = dct.__init_dctN(b.shape[0])
     return numpy.dot(numpy.transpose(c), numpy.dot(b, c))
+